run ruff
This commit is contained in:
mcDandy
@@ -1,2 +1 @@
|
||||
from .more_math.nodes import comfy_entrypoint
|
||||
|
||||
|
||||
+26
-20
@@ -1,11 +1,12 @@
|
||||
import torch
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like,as_tensor
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like, as_tensor
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
from .MathNodeBase import MathNodeBase
|
||||
|
||||
|
||||
class AudioMathNode(MathNodeBase):
|
||||
"""
|
||||
Enables math expressions on Audio tensors.
|
||||
@@ -43,31 +44,36 @@ class AudioMathNode(MathNodeBase):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, AudioExpr, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
av = a['waveform']
|
||||
sample_rate = a['sample_rate']
|
||||
av = a["waveform"]
|
||||
sample_rate = a["sample_rate"]
|
||||
|
||||
a, b, c, d = cls.prepare_inputs(a, b, c, d)
|
||||
|
||||
bv, cv, dv = b['waveform'], c['waveform'], d['waveform']
|
||||
bv, cv, dv = b["waveform"], c["waveform"], d["waveform"]
|
||||
|
||||
variables = {
|
||||
'a': av, 'b': bv, 'c': cv, 'd': dv,
|
||||
'w': w, 'x': x, 'y': y, 'z': z,
|
||||
'B': getIndexTensorAlongDim(av, 0),
|
||||
'C': getIndexTensorAlongDim(av, 1),
|
||||
'S': getIndexTensorAlongDim(av, 2),
|
||||
'R': torch.full_like(av, sample_rate, dtype=torch.float32),
|
||||
'T': torch.full_like(av, av.shape[2], dtype=torch.float32),
|
||||
'N': av.shape[1],
|
||||
'batch': getIndexTensorAlongDim(av, 0),
|
||||
'channel': getIndexTensorAlongDim(av, 1),
|
||||
'sample': getIndexTensorAlongDim(av, 2),
|
||||
'sample_rate': torch.full_like(av, sample_rate, dtype=torch.float32),
|
||||
'sample_count': torch.full_like(av, av.shape[2], dtype=torch.float32),
|
||||
'channel_count': av.shape[1],
|
||||
"a": av,
|
||||
"b": bv,
|
||||
"c": cv,
|
||||
"d": dv,
|
||||
"w": w,
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z": z,
|
||||
"B": getIndexTensorAlongDim(av, 0),
|
||||
"C": getIndexTensorAlongDim(av, 1),
|
||||
"S": getIndexTensorAlongDim(av, 2),
|
||||
"R": torch.full_like(av, sample_rate, dtype=torch.float32),
|
||||
"T": torch.full_like(av, av.shape[2], dtype=torch.float32),
|
||||
"N": av.shape[1],
|
||||
"batch": getIndexTensorAlongDim(av, 0),
|
||||
"channel": getIndexTensorAlongDim(av, 1),
|
||||
"sample": getIndexTensorAlongDim(av, 2),
|
||||
"sample_rate": torch.full_like(av, sample_rate, dtype=torch.float32),
|
||||
"sample_count": torch.full_like(av, av.shape[2], dtype=torch.float32),
|
||||
"channel_count": av.shape[1],
|
||||
} | generate_dim_variables(av)
|
||||
|
||||
|
||||
result_tensor = eval_tensor_expr(AudioExpr, variables, av.shape)
|
||||
|
||||
return ({'waveform': as_tensor(result_tensor,a.shape), 'sample_rate': sample_rate},)
|
||||
return ({"waveform": as_tensor(result_tensor, a.shape), "sample_rate": sample_rate},)
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
from comfy_api.latest import io
|
||||
import torch
|
||||
|
||||
windows = {'bartlet':torch.bartlett_window, 'blackman':torch.blackman_window, 'hamming':torch.hamming_window,'hann':torch.hann_window}
|
||||
windows = {"bartlet": torch.bartlett_window, "blackman": torch.blackman_window, "hamming": torch.hamming_window, "hann": torch.hann_window}
|
||||
|
||||
|
||||
class AudioToSpectrogram(io.ComfyNode):
|
||||
"""
|
||||
@@ -9,6 +10,7 @@ class AudioToSpectrogram(io.ComfyNode):
|
||||
Red = Real, Green = logarithm of value (just so it looks good), Blue = Imaginary.
|
||||
Each audio channel is stacked vertically.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -19,7 +21,9 @@ class AudioToSpectrogram(io.ComfyNode):
|
||||
io.Audio.Input(id="audio", tooltip="Input audio"),
|
||||
io.Int.Input(id="window_length", default=1024, min=16, max=4096, tooltip="Window length in samples"),
|
||||
io.Int.Input(id="hop_length", default=256, min=1, max=4096, tooltip="Stride of the window (hop length) in samples"),
|
||||
io.Int.Input(id="bucket_count", default=513, min=2, max=4096, tooltip="Number of frequency buckets (determines resolution)"),
|
||||
io.Int.Input(
|
||||
id="bucket_count", default=513, min=2, max=4096, tooltip="Number of frequency buckets (determines resolution)"
|
||||
),
|
||||
io.Combo.Input(id="window_type", default="hann", options=list(windows.keys()), tooltip="Type of window function to apply"),
|
||||
],
|
||||
outputs=[
|
||||
@@ -30,26 +34,29 @@ class AudioToSpectrogram(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, window_length, hop_length, bucket_count,window_type):
|
||||
waveform = audio['waveform']
|
||||
sample_rate = audio['sample_rate']
|
||||
def execute(cls, audio, window_length, hop_length, bucket_count, window_type):
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
B, C, S = waveform.shape
|
||||
n_fft = (bucket_count - 1) * 2
|
||||
flat_waveform = waveform.reshape(B * C, S)
|
||||
window = windows[window_type](window_length, device=waveform.device)
|
||||
|
||||
stft_out = torch.stft(
|
||||
flat_waveform, n_fft=n_fft, hop_length=hop_length, win_length=window_length,
|
||||
window=window, center=True, pad_mode='reflect', normalized=False,
|
||||
onesided=True, return_complex=True
|
||||
flat_waveform,
|
||||
n_fft=n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=window_length,
|
||||
window=window,
|
||||
center=True,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
|
||||
F, T = stft_out.shape[1], stft_out.shape[2]
|
||||
stft_out = stft_out.reshape(B, C, F, T).reshape(B, C * F, T)
|
||||
image = torch.stack([
|
||||
stft_out.real,
|
||||
torch.log1p(torch.abs(stft_out)),
|
||||
stft_out.imag
|
||||
], dim=-1)
|
||||
image = torch.stack([stft_out.real, torch.log1p(torch.abs(stft_out)), stft_out.imag], dim=-1)
|
||||
print(image.type())
|
||||
return (image, C, sample_rate)
|
||||
return (image, C, sample_rate)
|
||||
|
||||
@@ -3,10 +3,12 @@ from .modelLikeCommon import calculate_patches
|
||||
from inspect import cleandoc
|
||||
from .helper_functions import comonLazy
|
||||
|
||||
|
||||
class CLIPMathNode(io.ComfyNode):
|
||||
"""
|
||||
This node enables the use of math expressions on CLIP weights.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -15,13 +17,13 @@ class CLIPMathNode(io.ComfyNode):
|
||||
category="More math",
|
||||
inputs=[
|
||||
io.Clip.Input(id="a", tooltip="Main CLIP (base)"),
|
||||
io.Clip.Input(id="b", optional=True,lazy=True, tooltip="Optional 2nd CLIP"),
|
||||
io.Clip.Input(id="c", optional=True,lazy=True, tooltip="Optional 3rd CLIP"),
|
||||
io.Clip.Input(id="d", optional=True,lazy=True, tooltip="Optional 4th CLIP"),
|
||||
io.Float.Input(id="w", default=0.0, optional=True,lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True,lazy=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,lazy=True, force_input=True),
|
||||
io.Clip.Input(id="b", optional=True, lazy=True, tooltip="Optional 2nd CLIP"),
|
||||
io.Clip.Input(id="c", optional=True, lazy=True, tooltip="Optional 3rd CLIP"),
|
||||
io.Clip.Input(id="d", optional=True, lazy=True, tooltip="Optional 4th CLIP"),
|
||||
io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True, lazy=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, lazy=True, force_input=True),
|
||||
io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
|
||||
],
|
||||
outputs=[
|
||||
@@ -30,9 +32,10 @@ class CLIPMathNode(io.ComfyNode):
|
||||
)
|
||||
|
||||
tooltip = cleandoc(__doc__)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
|
||||
return comonLazy(Model, a, b, c, d,w,x,y,z)
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[], w=0, x=0, y=0, z=0):
|
||||
return comonLazy(Model, a, b, c, d, w, x, y, z)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
from inspect import cleandoc
|
||||
import torch
|
||||
|
||||
from .helper_functions import comonLazy, eval_tensor_expr, generate_dim_variables, make_zero_like,as_tensor
|
||||
from .helper_functions import comonLazy, eval_tensor_expr, generate_dim_variables, make_zero_like, as_tensor
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
|
||||
from .MathNodeBase import MathNodeBase
|
||||
|
||||
|
||||
class ConditioningMathNode(MathNodeBase):
|
||||
"""
|
||||
Enables math operations on conditionings.
|
||||
@@ -39,7 +39,9 @@ class ConditioningMathNode(MathNodeBase):
|
||||
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, lazy=True, force_input=True),
|
||||
io.String.Input(id="Tensor", default="a*(1-w)+b*w", tooltip="Expression for tensor part (image composition)"),
|
||||
io.String.Input(id="pooled_output", default="a*(1-w)+b*w", tooltip="Expression for pooled output (condensed representation)"),
|
||||
io.String.Input(
|
||||
id="pooled_output", default="a*(1-w)+b*w", tooltip="Expression for pooled output (condensed representation)"
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(),
|
||||
@@ -63,8 +65,8 @@ class ConditioningMathNode(MathNodeBase):
|
||||
ta, tb, tc, td = a[0][0], b[0][0], c[0][0], d[0][0]
|
||||
|
||||
# Evaluate tensor expression
|
||||
variables = {'a': ta, 'b': tb, 'c': tc, 'd': td, 'w': w, 'x': x, 'y': y, 'z': z} | generate_dim_variables(ta)
|
||||
result_tensor = as_tensor(eval_tensor_expr(Tensor, variables, ta.shape),ta.shape)
|
||||
variables = {"a": ta, "b": tb, "c": tc, "d": td, "w": w, "x": x, "y": y, "z": z} | generate_dim_variables(ta)
|
||||
result_tensor = as_tensor(eval_tensor_expr(Tensor, variables, ta.shape), ta.shape)
|
||||
|
||||
# Evaluate pooled_output expression if available
|
||||
pa = a[0][1].get("pooled_output")
|
||||
@@ -72,8 +74,8 @@ class ConditioningMathNode(MathNodeBase):
|
||||
pb = b[0][1].get("pooled_output")
|
||||
pc = c[0][1].get("pooled_output")
|
||||
pd = d[0][1].get("pooled_output")
|
||||
variables = {'a': pa, 'b': pb, 'c': pc, 'd': pd, 'w': w, 'x': x, 'y': y, 'z': z} | generate_dim_variables(pa)
|
||||
result_pooled = as_tensor(eval_tensor_expr(pooled_output, variables, pa.shape),pa.shape)
|
||||
variables = {"a": pa, "b": pb, "c": pc, "d": pd, "w": w, "x": x, "y": y, "z": z} | generate_dim_variables(pa)
|
||||
result_pooled = as_tensor(eval_tensor_expr(pooled_output, variables, pa.shape), pa.shape)
|
||||
else:
|
||||
result_pooled = None
|
||||
|
||||
|
||||
@@ -3,21 +3,23 @@ from inspect import cleandoc
|
||||
from .helper_functions import comonLazy, eval_float_expr
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
class FloatMathNode(io.ComfyNode):
|
||||
"""
|
||||
This node enables the use of math expressions on Floats.
|
||||
|
||||
|
||||
Inputs:
|
||||
a, b, c, d: Floats, bound to variables with the same name.
|
||||
w, x, y, z: Floats, bound to variables of the expression.
|
||||
FloatFunc: String, describing math expression.
|
||||
FloatFunc: String, describing math expression.
|
||||
Valid functions: sin, cos, tan, abs, sqrt, min, max, norm, etc.
|
||||
Operators: +, -, *, /, ^, %.
|
||||
Operators: +, -, *, /, ^, %.
|
||||
Constants: e, pi.
|
||||
|
||||
|
||||
Outputs:
|
||||
FLOAT: The result of evaluating the math expression.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -48,6 +50,6 @@ class FloatMathNode(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, FloatFunc, a, b=0.0, c=0.0, d=0.0, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
variables = {'a': a, 'b': b, 'c': c, 'd': d, 'w': w, 'x': x, 'y': y, 'z': z}
|
||||
variables = {"a": a, "b": b, "c": c, "d": d, "w": w, "x": x, "y": y, "z": z}
|
||||
result = eval_float_expr(FloatFunc, variables)
|
||||
return (result,)
|
||||
|
||||
+25
-13
@@ -1,11 +1,12 @@
|
||||
import torch
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like,as_tensor
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, eval_tensor_expr, make_zero_like, as_tensor
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
from .MathNodeBase import MathNodeBase
|
||||
|
||||
|
||||
class ImageMathNode(MathNodeBase):
|
||||
"""
|
||||
Enables math expressions on Images.
|
||||
@@ -45,20 +46,31 @@ class ImageMathNode(MathNodeBase):
|
||||
def execute(cls, Image, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
a, b, c, d = cls.prepare_inputs(a, b, c, d)
|
||||
|
||||
|
||||
variables = {
|
||||
'a': a, 'b': b, 'c': c, 'd': d,
|
||||
'w': w, 'x': x, 'y': y, 'z': z,
|
||||
'X': getIndexTensorAlongDim(a, 3),
|
||||
'Y': getIndexTensorAlongDim(a, 2),
|
||||
'B': getIndexTensorAlongDim(a, 0), 'batch': getIndexTensorAlongDim(a, 0),
|
||||
'C': getIndexTensorAlongDim(a, 1), 'channel': getIndexTensorAlongDim(a, 1),
|
||||
'W': a.shape[3], 'width': a.shape[3],
|
||||
'H': a.shape[2], 'height': a.shape[2],
|
||||
'T': a.shape[0], 'batch_count': a.shape[0],
|
||||
'N': a.shape[1], 'channel_count': a.shape[1],
|
||||
"a": a,
|
||||
"b": b,
|
||||
"c": c,
|
||||
"d": d,
|
||||
"w": w,
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z": z,
|
||||
"X": getIndexTensorAlongDim(a, 3),
|
||||
"Y": getIndexTensorAlongDim(a, 2),
|
||||
"B": getIndexTensorAlongDim(a, 0),
|
||||
"batch": getIndexTensorAlongDim(a, 0),
|
||||
"C": getIndexTensorAlongDim(a, 1),
|
||||
"channel": getIndexTensorAlongDim(a, 1),
|
||||
"W": a.shape[3],
|
||||
"width": a.shape[3],
|
||||
"H": a.shape[2],
|
||||
"height": a.shape[2],
|
||||
"T": a.shape[0],
|
||||
"batch_count": a.shape[0],
|
||||
"N": a.shape[1],
|
||||
"channel_count": a.shape[1],
|
||||
} | generate_dim_variables(a)
|
||||
|
||||
result = eval_tensor_expr(Image, variables, a.shape)
|
||||
|
||||
return (as_tensor(result,a.shape),)
|
||||
return (as_tensor(result, a.shape),)
|
||||
|
||||
+47
-28
@@ -2,11 +2,18 @@ from inspect import cleandoc
|
||||
import comfy.nested_tensor
|
||||
from comfy_api.latest import io
|
||||
import torch
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like,as_tensor
|
||||
from .helper_functions import (
|
||||
generate_dim_variables,
|
||||
getIndexTensorAlongDim,
|
||||
comonLazy,
|
||||
parse_expr,
|
||||
eval_tensor_expr_with_tree,
|
||||
make_zero_like,
|
||||
as_tensor,
|
||||
)
|
||||
from .MathNodeBase import MathNodeBase
|
||||
|
||||
|
||||
|
||||
class LatentMathNode(MathNodeBase):
|
||||
"""
|
||||
This node enables the use of math expressions on Latents.
|
||||
@@ -22,13 +29,13 @@ class LatentMathNode(MathNodeBase):
|
||||
LATENT:
|
||||
Returns a LATENT object that contains the result of the math expression applied to the input conditionings.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
"""
|
||||
"""
|
||||
""" """
|
||||
return io.Schema(
|
||||
node_id="mrmth_LatentMathNode",
|
||||
display_name="Latent math",
|
||||
@@ -38,10 +45,10 @@ class LatentMathNode(MathNodeBase):
|
||||
io.Latent.Input(id="b", optional=True, lazy=True),
|
||||
io.Latent.Input(id="c", optional=True, lazy=True),
|
||||
io.Latent.Input(id="d", optional=True, lazy=True),
|
||||
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0,optional=True,lazy=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,lazy=True, force_input=True),
|
||||
io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True, lazy=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, lazy=True, force_input=True),
|
||||
io.String.Input(id="Latent", default="a*(1-w)+b*w", tooltip="Expression to apply on input latents"),
|
||||
],
|
||||
outputs=[
|
||||
@@ -49,11 +56,11 @@ class LatentMathNode(MathNodeBase):
|
||||
],
|
||||
)
|
||||
|
||||
#RETURN_NAMES = ("image_output_name",)
|
||||
# RETURN_NAMES = ("image_output_name",)
|
||||
tooltip = cleandoc(__doc__)
|
||||
|
||||
#OUTPUT_NODE = False
|
||||
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
||||
# OUTPUT_NODE = False
|
||||
# OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
||||
@classmethod
|
||||
def execute(cls, Latent, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
|
||||
# Extract raw sample tensors (may be Tensor or NestedTensor)
|
||||
@@ -87,25 +94,37 @@ class LatentMathNode(MathNodeBase):
|
||||
frame_count = a_t.shape[time_dim] if time_dim is not None else a_t.shape[batch_dim]
|
||||
|
||||
variables = {
|
||||
'a': a_t, 'b': b_t, 'c': c_t, 'd': d_t,
|
||||
'w': w, 'x': x, 'y': y, 'z': z,
|
||||
|
||||
'X': W, 'Y': H,
|
||||
'B': B, 'batch': B,
|
||||
'C': C, 'channel': C,
|
||||
'W': width_val, 'width': width_val,
|
||||
'H': height_val, 'height': height_val,
|
||||
'T': frame_count, 'batch_count': batch_count,
|
||||
'N': channel_count, 'channel_count': channel_count,
|
||||
"a": a_t,
|
||||
"b": b_t,
|
||||
"c": c_t,
|
||||
"d": d_t,
|
||||
"w": w,
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z": z,
|
||||
"X": W,
|
||||
"Y": H,
|
||||
"B": B,
|
||||
"batch": B,
|
||||
"C": C,
|
||||
"channel": C,
|
||||
"W": width_val,
|
||||
"width": width_val,
|
||||
"H": height_val,
|
||||
"height": height_val,
|
||||
"T": frame_count,
|
||||
"batch_count": batch_count,
|
||||
"N": channel_count,
|
||||
"channel_count": channel_count,
|
||||
} | generate_dim_variables(a_t)
|
||||
|
||||
if time_dim is not None:
|
||||
F = getIndexTensorAlongDim(a_t, time_dim)
|
||||
variables.update({'frame_idx': F, 'frame': F, 'frame_count': frame_count})
|
||||
variables.update({"frame_idx": F, "frame": F, "frame_count": frame_count})
|
||||
|
||||
return as_tensor(eval_tensor_expr_with_tree(tree, variables, a_t.shape),a_t.shape)
|
||||
return as_tensor(eval_tensor_expr_with_tree(tree, variables, a_t.shape), a_t.shape)
|
||||
|
||||
if hasattr(a_in, 'is_nested') and getattr(a_in, 'is_nested'):
|
||||
if hasattr(a_in, "is_nested") and getattr(a_in, "is_nested"):
|
||||
a_list = a_in.unbind()
|
||||
sizes = [t.shape[0] for t in a_list]
|
||||
merged_a = torch.cat(a_list, dim=0)
|
||||
@@ -113,7 +132,7 @@ class LatentMathNode(MathNodeBase):
|
||||
def merge_to_tensor(val, ref):
|
||||
if val is None:
|
||||
return make_zero_like(ref)
|
||||
if hasattr(val, 'is_nested') and getattr(val, 'is_nested'):
|
||||
if hasattr(val, "is_nested") and getattr(val, "is_nested"):
|
||||
lst = val.unbind()
|
||||
return torch.cat(lst, dim=0)
|
||||
if isinstance(val, (list, tuple)):
|
||||
@@ -143,7 +162,7 @@ class LatentMathNode(MathNodeBase):
|
||||
def to_tensor(val, ref):
|
||||
if val is None:
|
||||
return make_zero_like(ref)
|
||||
if hasattr(val, 'is_nested') and getattr(val, 'is_nested'):
|
||||
if hasattr(val, "is_nested") and getattr(val, "is_nested"):
|
||||
lst = val.unbind()
|
||||
elif isinstance(val, (list, tuple)):
|
||||
lst = list(val)
|
||||
@@ -184,6 +203,6 @@ class LatentMathNode(MathNodeBase):
|
||||
This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash
|
||||
changes between executions the LoadImage node is executed again.
|
||||
"""
|
||||
#@classmethod
|
||||
#def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
||||
# @classmethod
|
||||
# def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
||||
# return ""
|
||||
|
||||
+17
-16
@@ -1,6 +1,7 @@
|
||||
from comfy_api.latest import io
|
||||
from .helper_functions import comonLazy, make_zero_like
|
||||
|
||||
|
||||
class MathNodeBase(io.ComfyNode):
|
||||
"""
|
||||
Base class for More Math nodes providing common utilities.
|
||||
@@ -13,35 +14,35 @@ class MathNodeBase(io.ComfyNode):
|
||||
Finds the first argument name that isn't a known variable (a-d, w-z) and assumes it's the expression.
|
||||
"""
|
||||
# Known variable inputs (and internal ones we might ignore)
|
||||
known_vars = {'a', 'b', 'c', 'd', 'w', 'x', 'y', 'z', 'pooled_output'}
|
||||
|
||||
known_vars = {"a", "b", "c", "d", "w", "x", "y", "z", "pooled_output"}
|
||||
|
||||
# Heuristic: The expression is usually the first argument that is not a known variable.
|
||||
# But `comonLazy` takes (expr, a, b, c, d, w, x, y, z).
|
||||
# We need to extract the values from kwargs correctly.
|
||||
|
||||
|
||||
expr_key = None
|
||||
# Try finding the expression key
|
||||
for k in kwargs:
|
||||
if k not in known_vars and isinstance(kwargs[k], str):
|
||||
expr_key = k
|
||||
break
|
||||
|
||||
|
||||
# If we can't find it heuristically (e.g. standard naming), fallback or error?
|
||||
# Let's try to be safe. If we can't find it, we might be in a node with custom logic like ConditioningMathNode.
|
||||
# ConditioningMathNode has 'Tensor' and 'pooled_output' as expressions.
|
||||
|
||||
|
||||
if expr_key:
|
||||
expr = kwargs[expr_key]
|
||||
a = kwargs.get('a', None)
|
||||
b = kwargs.get('b', [])
|
||||
c = kwargs.get('c', [])
|
||||
d = kwargs.get('d', [])
|
||||
w = kwargs.get('w', 0)
|
||||
x = kwargs.get('x', 0)
|
||||
y = kwargs.get('y', 0)
|
||||
z = kwargs.get('z', 0)
|
||||
return comonLazy(expr, a, b, c, d, w, x, y, z)
|
||||
|
||||
expr = kwargs[expr_key]
|
||||
a = kwargs.get("a", None)
|
||||
b = kwargs.get("b", [])
|
||||
c = kwargs.get("c", [])
|
||||
d = kwargs.get("d", [])
|
||||
w = kwargs.get("w", 0)
|
||||
x = kwargs.get("x", 0)
|
||||
y = kwargs.get("y", 0)
|
||||
z = kwargs.get("z", 0)
|
||||
return comonLazy(expr, a, b, c, d, w, x, y, z)
|
||||
|
||||
# Fallback/Empty return if we can't determine dependencies automatically
|
||||
return []
|
||||
|
||||
|
||||
+11
-11
@@ -4,12 +4,12 @@ from .helper_functions import comonLazy
|
||||
from .modelLikeCommon import calculate_patches
|
||||
|
||||
|
||||
|
||||
class ModelMathNode(io.ComfyNode):
|
||||
"""
|
||||
This node enables the use of math expressions on Model weights (state_dict).
|
||||
Functionally acts as a custom model merge.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -18,13 +18,13 @@ class ModelMathNode(io.ComfyNode):
|
||||
category="More math",
|
||||
inputs=[
|
||||
io.Model.Input(id="a", tooltip="Main model (base)"),
|
||||
io.Model.Input(id="b", optional=True,lazy=True, tooltip="Optional 2nd model"),
|
||||
io.Model.Input(id="c", optional=True,lazy=True, tooltip="Optional 3rd model"),
|
||||
io.Model.Input(id="d", optional=True,lazy=True, tooltip="Optional 4th model"),
|
||||
io.Float.Input(id="w", default=0.0, optional=True,lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True,lazy=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,lazy=True, force_input=True),
|
||||
io.Model.Input(id="b", optional=True, lazy=True, tooltip="Optional 2nd model"),
|
||||
io.Model.Input(id="c", optional=True, lazy=True, tooltip="Optional 3rd model"),
|
||||
io.Model.Input(id="d", optional=True, lazy=True, tooltip="Optional 4th model"),
|
||||
io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True, lazy=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, lazy=True, force_input=True),
|
||||
io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
|
||||
],
|
||||
outputs=[
|
||||
@@ -35,8 +35,8 @@ class ModelMathNode(io.ComfyNode):
|
||||
tooltip = cleandoc(__doc__)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
|
||||
return comonLazy(Model, a, b, c, d,w,x,y,z)
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[], w=0, x=0, y=0, z=0):
|
||||
return comonLazy(Model, a, b, c, d, w, x, y, z)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
|
||||
@@ -44,4 +44,4 @@ class ModelMathNode(io.ComfyNode):
|
||||
out_model = a.clone()
|
||||
if patches:
|
||||
out_model.add_patches(patches, 1.0, 1.0)
|
||||
return (out_model,)
|
||||
return (out_model,)
|
||||
|
||||
+70
-39
@@ -2,14 +2,21 @@ from inspect import cleandoc
|
||||
|
||||
import torch
|
||||
|
||||
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, comonLazy, parse_expr, eval_tensor_expr_with_tree, make_zero_like,as_tensor
|
||||
from .helper_functions import (
|
||||
generate_dim_variables,
|
||||
getIndexTensorAlongDim,
|
||||
comonLazy,
|
||||
parse_expr,
|
||||
eval_tensor_expr_with_tree,
|
||||
make_zero_like,
|
||||
as_tensor,
|
||||
)
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
import comfy.nested_tensor as _nested_tensor_module
|
||||
|
||||
|
||||
|
||||
class NoiseMathNode(io.ComfyNode):
|
||||
"""
|
||||
This node enables the use of math expressions on Latents.
|
||||
@@ -25,40 +32,44 @@ class NoiseMathNode(io.ComfyNode):
|
||||
LATENT:
|
||||
Returns a LATENT object that contains the result of the math expression applied to the input conditionings.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
"""
|
||||
"""
|
||||
""" """
|
||||
return io.Schema(
|
||||
node_id="mrmth_NoiseMathNode",
|
||||
category="More math",
|
||||
display_name="Noise math",
|
||||
inputs=[
|
||||
io.Noise.Input(id="a"),
|
||||
io.Noise.Input(id="b", optional=True,lazy=True),
|
||||
io.Noise.Input(id="c", optional=True,lazy=True),
|
||||
io.Noise.Input(id="d", optional=True,lazy=True),
|
||||
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0,optional=True,lazy=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,lazy=True, force_input=True),
|
||||
io.Noise.Input(id="b", optional=True, lazy=True),
|
||||
io.Noise.Input(id="c", optional=True, lazy=True),
|
||||
io.Noise.Input(id="d", optional=True, lazy=True),
|
||||
io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, optional=True, lazy=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, lazy=True, force_input=True),
|
||||
io.String.Input(id="Noise", default="a*(1-w)+b*w", tooltip="Expression to apply on input noise generators"),
|
||||
],
|
||||
outputs=[
|
||||
io.Noise.Output(),
|
||||
],
|
||||
)
|
||||
#RETURN_NAMES = ("image_output_name",)
|
||||
|
||||
# RETURN_NAMES = ("image_output_name",)
|
||||
tooltip = cleandoc(__doc__)
|
||||
#OUTPUT_NODE = False
|
||||
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
||||
# OUTPUT_NODE = False
|
||||
# OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
||||
|
||||
CATEGORY = "More math"
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, Noise, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
|
||||
return comonLazy(Noise, a, b, c, d,w,x,y,z)
|
||||
def check_lazy_status(cls, Noise, a, b=[], c=[], d=[], w=0, x=0, y=0, z=0):
|
||||
return comonLazy(Noise, a, b, c, d, w, x, y, z)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, Noise, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
return (NoiseExecutor(a, b, c, d, w, x, y, z, Noise),)
|
||||
@@ -71,15 +82,13 @@ class NoiseMathNode(io.ComfyNode):
|
||||
This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash
|
||||
changes between executions the LoadImage node is executed again.
|
||||
"""
|
||||
#@classmethod
|
||||
#def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
||||
# @classmethod
|
||||
# def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
||||
# return ""
|
||||
|
||||
|
||||
|
||||
|
||||
class NoiseExecutor():
|
||||
def __init__(self,a,b,c,d,w,x,y,z, Noise):
|
||||
class NoiseExecutor:
|
||||
def __init__(self, a, b, c, d, w, x, y, z, Noise):
|
||||
self.a = a
|
||||
self.b = b
|
||||
self.c = c
|
||||
@@ -90,8 +99,10 @@ class NoiseExecutor():
|
||||
self.z = z
|
||||
self.expr = Noise
|
||||
self.tree = parse_expr(Noise)
|
||||
seed = -1;
|
||||
def generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor:
|
||||
|
||||
seed = -1
|
||||
|
||||
def generate_noise(self, input_latent: torch.Tensor) -> torch.Tensor:
|
||||
samples = input_latent["samples"]
|
||||
|
||||
a_val = self.a.generate_noise(input_latent) if self.a is not None else make_zero_like(samples)
|
||||
@@ -104,7 +115,7 @@ class NoiseExecutor():
|
||||
if val is None:
|
||||
return [make_zero_like(r) for r in ref_list]
|
||||
# If val is a NestedTensor-like, return underlying list
|
||||
if hasattr(val, 'is_nested') and getattr(val, 'is_nested'):
|
||||
if hasattr(val, "is_nested") and getattr(val, "is_nested"):
|
||||
return val.unbind()
|
||||
if isinstance(val, list) or isinstance(val, tuple):
|
||||
return list(val)
|
||||
@@ -121,7 +132,7 @@ class NoiseExecutor():
|
||||
return [val for _ in ref_list]
|
||||
|
||||
# nested case: merge subtensors, evaluate once, split back
|
||||
if hasattr(samples, 'is_nested') and getattr(samples, 'is_nested'):
|
||||
if hasattr(samples, "is_nested") and getattr(samples, "is_nested"):
|
||||
sample_list = samples.unbind()
|
||||
sizes = [t.shape[0] for t in sample_list]
|
||||
merged_samples = torch.cat(sample_list, dim=0)
|
||||
@@ -129,7 +140,7 @@ class NoiseExecutor():
|
||||
def merge_to_tensor(val, ref):
|
||||
if val is None:
|
||||
return make_zero_like(ref)
|
||||
if hasattr(val, 'is_nested') and getattr(val, 'is_nested'):
|
||||
if hasattr(val, "is_nested") and getattr(val, "is_nested"):
|
||||
lst = val.unbind()
|
||||
return torch.cat(lst, dim=0)
|
||||
if isinstance(val, (list, tuple)):
|
||||
@@ -142,7 +153,10 @@ class NoiseExecutor():
|
||||
# if val has per-subtensor batches, try to split and cat
|
||||
try:
|
||||
if val.shape[0] == len(sample_list):
|
||||
return torch.cat([val[i].unsqueeze(0).expand(sample_list[i].shape[0], *val.shape[1:]) for i in range(len(sample_list))], dim=0)
|
||||
return torch.cat(
|
||||
[val[i].unsqueeze(0).expand(sample_list[i].shape[0], *val.shape[1:]) for i in range(len(sample_list))],
|
||||
dim=0,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return make_zero_like(ref)
|
||||
@@ -176,21 +190,38 @@ class NoiseExecutor():
|
||||
H = getIndexTensorAlongDim(merged_samples, height_dim)
|
||||
C = getIndexTensorAlongDim(merged_samples, channel_dim)
|
||||
variables = {
|
||||
'a': merged_a, 'b': merged_b, 'c': merged_c, 'd': merged_d,
|
||||
'w': self.w, 'x': self.x, 'y': self.y, 'z': self.z,
|
||||
'B': B, 'X': W, 'Y': H, 'C': C,
|
||||
'W': merged_samples.shape[width_dim], 'H': merged_samples.shape[height_dim],
|
||||
'I': merged_samples, 'T': frame_count, 'N': merged_samples.shape[channel_dim],
|
||||
'batch': B, 'width': merged_samples.shape[width_dim], 'height': merged_samples.shape[height_dim], 'channel': C,
|
||||
'batch_count': merged_samples.shape[0], 'channel_count': merged_samples.shape[1], 'input_latent': merged_samples,
|
||||
"a": merged_a,
|
||||
"b": merged_b,
|
||||
"c": merged_c,
|
||||
"d": merged_d,
|
||||
"w": self.w,
|
||||
"x": self.x,
|
||||
"y": self.y,
|
||||
"z": self.z,
|
||||
"B": B,
|
||||
"X": W,
|
||||
"Y": H,
|
||||
"C": C,
|
||||
"W": merged_samples.shape[width_dim],
|
||||
"H": merged_samples.shape[height_dim],
|
||||
"I": merged_samples,
|
||||
"T": frame_count,
|
||||
"N": merged_samples.shape[channel_dim],
|
||||
"batch": B,
|
||||
"width": merged_samples.shape[width_dim],
|
||||
"height": merged_samples.shape[height_dim],
|
||||
"channel": C,
|
||||
"batch_count": merged_samples.shape[0],
|
||||
"channel_count": merged_samples.shape[1],
|
||||
"input_latent": merged_samples,
|
||||
} | generate_dim_variables(merged_samples)
|
||||
if time_dim is not None:
|
||||
F = getIndexTensorAlongDim(merged_samples, time_dim)
|
||||
variables.update({'frame': F, 'frame_count': frame_count})
|
||||
variables.update({"frame": F, "frame_count": frame_count})
|
||||
|
||||
merged_result = as_tensor(eval_tensor_expr_with_tree(self.tree, variables, variables['a'].shape),variables['a'].shape)
|
||||
merged_result = as_tensor(eval_tensor_expr_with_tree(self.tree, variables, variables["a"].shape), variables["a"].shape)
|
||||
|
||||
if hasattr(samples, 'is_nested') and getattr(samples, 'is_nested'):
|
||||
if hasattr(samples, "is_nested") and getattr(samples, "is_nested"):
|
||||
split_results = list(merged_result.split(sizes, dim=0))
|
||||
return _nested_tensor_module.NestedTensor(split_results)
|
||||
|
||||
|
||||
+4491
-216
File diff suppressed because it is too large
Load Diff
@@ -1,741 +1,661 @@
|
||||
# Generated from ./MathExpr.g4 by ANTLR 4.13.2
|
||||
from antlr4 import *
|
||||
|
||||
if "." in __name__:
|
||||
from .MathExprParser import MathExprParser
|
||||
else:
|
||||
from MathExprParser import MathExprParser
|
||||
|
||||
|
||||
# This class defines a complete listener for a parse tree produced by MathExprParser.
|
||||
class MathExprListener(ParseTreeListener):
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#expr.
|
||||
def enterExpr(self, ctx:MathExprParser.ExprContext):
|
||||
def enterExpr(self, ctx: MathExprParser.ExprContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#expr.
|
||||
def exitExpr(self, ctx:MathExprParser.ExprContext):
|
||||
def exitExpr(self, ctx: MathExprParser.ExprContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#LtExp.
|
||||
def enterLtExp(self, ctx:MathExprParser.LtExpContext):
|
||||
def enterLtExp(self, ctx: MathExprParser.LtExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#LtExp.
|
||||
def exitLtExp(self, ctx:MathExprParser.LtExpContext):
|
||||
def exitLtExp(self, ctx: MathExprParser.LtExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#EqExp.
|
||||
def enterEqExp(self, ctx:MathExprParser.EqExpContext):
|
||||
def enterEqExp(self, ctx: MathExprParser.EqExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#EqExp.
|
||||
def exitEqExp(self, ctx:MathExprParser.EqExpContext):
|
||||
def exitEqExp(self, ctx: MathExprParser.EqExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ToAdd.
|
||||
def enterToAdd(self, ctx:MathExprParser.ToAddContext):
|
||||
def enterToAdd(self, ctx: MathExprParser.ToAddContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ToAdd.
|
||||
def exitToAdd(self, ctx:MathExprParser.ToAddContext):
|
||||
def exitToAdd(self, ctx: MathExprParser.ToAddContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#GeExp.
|
||||
def enterGeExp(self, ctx:MathExprParser.GeExpContext):
|
||||
def enterGeExp(self, ctx: MathExprParser.GeExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#GeExp.
|
||||
def exitGeExp(self, ctx:MathExprParser.GeExpContext):
|
||||
def exitGeExp(self, ctx: MathExprParser.GeExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#LeExp.
|
||||
def enterLeExp(self, ctx:MathExprParser.LeExpContext):
|
||||
def enterLeExp(self, ctx: MathExprParser.LeExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#LeExp.
|
||||
def exitLeExp(self, ctx:MathExprParser.LeExpContext):
|
||||
def exitLeExp(self, ctx: MathExprParser.LeExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#NeExp.
|
||||
def enterNeExp(self, ctx:MathExprParser.NeExpContext):
|
||||
def enterNeExp(self, ctx: MathExprParser.NeExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#NeExp.
|
||||
def exitNeExp(self, ctx:MathExprParser.NeExpContext):
|
||||
def exitNeExp(self, ctx: MathExprParser.NeExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#GtExp.
|
||||
def enterGtExp(self, ctx:MathExprParser.GtExpContext):
|
||||
def enterGtExp(self, ctx: MathExprParser.GtExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#GtExp.
|
||||
def exitGtExp(self, ctx:MathExprParser.GtExpContext):
|
||||
def exitGtExp(self, ctx: MathExprParser.GtExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AddExp.
|
||||
def enterAddExp(self, ctx:MathExprParser.AddExpContext):
|
||||
def enterAddExp(self, ctx: MathExprParser.AddExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AddExp.
|
||||
def exitAddExp(self, ctx:MathExprParser.AddExpContext):
|
||||
def exitAddExp(self, ctx: MathExprParser.AddExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ToMul.
|
||||
def enterToMul(self, ctx:MathExprParser.ToMulContext):
|
||||
def enterToMul(self, ctx: MathExprParser.ToMulContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ToMul.
|
||||
def exitToMul(self, ctx:MathExprParser.ToMulContext):
|
||||
def exitToMul(self, ctx: MathExprParser.ToMulContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SubExp.
|
||||
def enterSubExp(self, ctx:MathExprParser.SubExpContext):
|
||||
def enterSubExp(self, ctx: MathExprParser.SubExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SubExp.
|
||||
def exitSubExp(self, ctx:MathExprParser.SubExpContext):
|
||||
def exitSubExp(self, ctx: MathExprParser.SubExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#MulExp.
|
||||
def enterMulExp(self, ctx:MathExprParser.MulExpContext):
|
||||
def enterMulExp(self, ctx: MathExprParser.MulExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#MulExp.
|
||||
def exitMulExp(self, ctx:MathExprParser.MulExpContext):
|
||||
def exitMulExp(self, ctx: MathExprParser.MulExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ModExp.
|
||||
def enterModExp(self, ctx:MathExprParser.ModExpContext):
|
||||
def enterModExp(self, ctx: MathExprParser.ModExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ModExp.
|
||||
def exitModExp(self, ctx:MathExprParser.ModExpContext):
|
||||
def exitModExp(self, ctx: MathExprParser.ModExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#DivExp.
|
||||
def enterDivExp(self, ctx:MathExprParser.DivExpContext):
|
||||
def enterDivExp(self, ctx: MathExprParser.DivExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#DivExp.
|
||||
def exitDivExp(self, ctx:MathExprParser.DivExpContext):
|
||||
def exitDivExp(self, ctx: MathExprParser.DivExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ToPow.
|
||||
def enterToPow(self, ctx:MathExprParser.ToPowContext):
|
||||
def enterToPow(self, ctx: MathExprParser.ToPowContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ToPow.
|
||||
def exitToPow(self, ctx:MathExprParser.ToPowContext):
|
||||
def exitToPow(self, ctx: MathExprParser.ToPowContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#PowExp.
|
||||
def enterPowExp(self, ctx:MathExprParser.PowExpContext):
|
||||
def enterPowExp(self, ctx: MathExprParser.PowExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#PowExp.
|
||||
def exitPowExp(self, ctx:MathExprParser.PowExpContext):
|
||||
def exitPowExp(self, ctx: MathExprParser.PowExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ToUnary.
|
||||
def enterToUnary(self, ctx:MathExprParser.ToUnaryContext):
|
||||
def enterToUnary(self, ctx: MathExprParser.ToUnaryContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ToUnary.
|
||||
def exitToUnary(self, ctx:MathExprParser.ToUnaryContext):
|
||||
def exitToUnary(self, ctx: MathExprParser.ToUnaryContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#UnaryPlus.
|
||||
def enterUnaryPlus(self, ctx:MathExprParser.UnaryPlusContext):
|
||||
def enterUnaryPlus(self, ctx: MathExprParser.UnaryPlusContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#UnaryPlus.
|
||||
def exitUnaryPlus(self, ctx:MathExprParser.UnaryPlusContext):
|
||||
def exitUnaryPlus(self, ctx: MathExprParser.UnaryPlusContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#UnaryMinus.
|
||||
def enterUnaryMinus(self, ctx:MathExprParser.UnaryMinusContext):
|
||||
def enterUnaryMinus(self, ctx: MathExprParser.UnaryMinusContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#UnaryMinus.
|
||||
def exitUnaryMinus(self, ctx:MathExprParser.UnaryMinusContext):
|
||||
def exitUnaryMinus(self, ctx: MathExprParser.UnaryMinusContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ToAtom.
|
||||
def enterToAtom(self, ctx:MathExprParser.ToAtomContext):
|
||||
def enterToAtom(self, ctx: MathExprParser.ToAtomContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ToAtom.
|
||||
def exitToAtom(self, ctx:MathExprParser.ToAtomContext):
|
||||
def exitToAtom(self, ctx: MathExprParser.ToAtomContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func1Exp.
|
||||
def enterFunc1Exp(self, ctx:MathExprParser.Func1ExpContext):
|
||||
def enterFunc1Exp(self, ctx: MathExprParser.Func1ExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Func1Exp.
|
||||
def exitFunc1Exp(self, ctx:MathExprParser.Func1ExpContext):
|
||||
def exitFunc1Exp(self, ctx: MathExprParser.Func1ExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func2Exp.
|
||||
def enterFunc2Exp(self, ctx:MathExprParser.Func2ExpContext):
|
||||
def enterFunc2Exp(self, ctx: MathExprParser.Func2ExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Func2Exp.
|
||||
def exitFunc2Exp(self, ctx:MathExprParser.Func2ExpContext):
|
||||
def exitFunc2Exp(self, ctx: MathExprParser.Func2ExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func3Exp.
|
||||
def enterFunc3Exp(self, ctx:MathExprParser.Func3ExpContext):
|
||||
def enterFunc3Exp(self, ctx: MathExprParser.Func3ExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Func3Exp.
|
||||
def exitFunc3Exp(self, ctx:MathExprParser.Func3ExpContext):
|
||||
def exitFunc3Exp(self, ctx: MathExprParser.Func3ExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Func4Exp.
|
||||
def enterFunc4Exp(self, ctx:MathExprParser.Func4ExpContext):
|
||||
def enterFunc4Exp(self, ctx: MathExprParser.Func4ExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Func4Exp.
|
||||
def exitFunc4Exp(self, ctx:MathExprParser.Func4ExpContext):
|
||||
def exitFunc4Exp(self, ctx: MathExprParser.Func4ExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#FuncNExp.
|
||||
def enterFuncNExp(self, ctx:MathExprParser.FuncNExpContext):
|
||||
def enterFuncNExp(self, ctx: MathExprParser.FuncNExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#FuncNExp.
|
||||
def exitFuncNExp(self, ctx:MathExprParser.FuncNExpContext):
|
||||
def exitFuncNExp(self, ctx: MathExprParser.FuncNExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#VariableExp.
|
||||
def enterVariableExp(self, ctx:MathExprParser.VariableExpContext):
|
||||
def enterVariableExp(self, ctx: MathExprParser.VariableExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#VariableExp.
|
||||
def exitVariableExp(self, ctx:MathExprParser.VariableExpContext):
|
||||
def exitVariableExp(self, ctx: MathExprParser.VariableExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#NumberExp.
|
||||
def enterNumberExp(self, ctx:MathExprParser.NumberExpContext):
|
||||
def enterNumberExp(self, ctx: MathExprParser.NumberExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#NumberExp.
|
||||
def exitNumberExp(self, ctx:MathExprParser.NumberExpContext):
|
||||
def exitNumberExp(self, ctx: MathExprParser.NumberExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ConstantExp.
|
||||
def enterConstantExp(self, ctx:MathExprParser.ConstantExpContext):
|
||||
def enterConstantExp(self, ctx: MathExprParser.ConstantExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ConstantExp.
|
||||
def exitConstantExp(self, ctx:MathExprParser.ConstantExpContext):
|
||||
def exitConstantExp(self, ctx: MathExprParser.ConstantExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ParenExp.
|
||||
def enterParenExp(self, ctx:MathExprParser.ParenExpContext):
|
||||
def enterParenExp(self, ctx: MathExprParser.ParenExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ParenExp.
|
||||
def exitParenExp(self, ctx:MathExprParser.ParenExpContext):
|
||||
def exitParenExp(self, ctx: MathExprParser.ParenExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AbsExp.
|
||||
def enterAbsExp(self, ctx:MathExprParser.AbsExpContext):
|
||||
def enterAbsExp(self, ctx: MathExprParser.AbsExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AbsExp.
|
||||
def exitAbsExp(self, ctx:MathExprParser.AbsExpContext):
|
||||
def exitAbsExp(self, ctx: MathExprParser.AbsExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ListExp.
|
||||
def enterListExp(self, ctx:MathExprParser.ListExpContext):
|
||||
def enterListExp(self, ctx: MathExprParser.ListExpContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ListExp.
|
||||
def exitListExp(self, ctx:MathExprParser.ListExpContext):
|
||||
def exitListExp(self, ctx: MathExprParser.ListExpContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SinFunc.
|
||||
def enterSinFunc(self, ctx:MathExprParser.SinFuncContext):
|
||||
def enterSinFunc(self, ctx: MathExprParser.SinFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SinFunc.
|
||||
def exitSinFunc(self, ctx:MathExprParser.SinFuncContext):
|
||||
def exitSinFunc(self, ctx: MathExprParser.SinFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#CosFunc.
|
||||
def enterCosFunc(self, ctx:MathExprParser.CosFuncContext):
|
||||
def enterCosFunc(self, ctx: MathExprParser.CosFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#CosFunc.
|
||||
def exitCosFunc(self, ctx:MathExprParser.CosFuncContext):
|
||||
def exitCosFunc(self, ctx: MathExprParser.CosFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#TanFunc.
|
||||
def enterTanFunc(self, ctx:MathExprParser.TanFuncContext):
|
||||
def enterTanFunc(self, ctx: MathExprParser.TanFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#TanFunc.
|
||||
def exitTanFunc(self, ctx:MathExprParser.TanFuncContext):
|
||||
def exitTanFunc(self, ctx: MathExprParser.TanFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AsinFunc.
|
||||
def enterAsinFunc(self, ctx:MathExprParser.AsinFuncContext):
|
||||
def enterAsinFunc(self, ctx: MathExprParser.AsinFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AsinFunc.
|
||||
def exitAsinFunc(self, ctx:MathExprParser.AsinFuncContext):
|
||||
def exitAsinFunc(self, ctx: MathExprParser.AsinFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AcosFunc.
|
||||
def enterAcosFunc(self, ctx:MathExprParser.AcosFuncContext):
|
||||
def enterAcosFunc(self, ctx: MathExprParser.AcosFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AcosFunc.
|
||||
def exitAcosFunc(self, ctx:MathExprParser.AcosFuncContext):
|
||||
def exitAcosFunc(self, ctx: MathExprParser.AcosFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AtanFunc.
|
||||
def enterAtanFunc(self, ctx:MathExprParser.AtanFuncContext):
|
||||
def enterAtanFunc(self, ctx: MathExprParser.AtanFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AtanFunc.
|
||||
def exitAtanFunc(self, ctx:MathExprParser.AtanFuncContext):
|
||||
def exitAtanFunc(self, ctx: MathExprParser.AtanFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SinhFunc.
|
||||
def enterSinhFunc(self, ctx:MathExprParser.SinhFuncContext):
|
||||
def enterSinhFunc(self, ctx: MathExprParser.SinhFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SinhFunc.
|
||||
def exitSinhFunc(self, ctx:MathExprParser.SinhFuncContext):
|
||||
def exitSinhFunc(self, ctx: MathExprParser.SinhFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#CoshFunc.
|
||||
def enterCoshFunc(self, ctx:MathExprParser.CoshFuncContext):
|
||||
def enterCoshFunc(self, ctx: MathExprParser.CoshFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#CoshFunc.
|
||||
def exitCoshFunc(self, ctx:MathExprParser.CoshFuncContext):
|
||||
def exitCoshFunc(self, ctx: MathExprParser.CoshFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#TanhFunc.
|
||||
def enterTanhFunc(self, ctx:MathExprParser.TanhFuncContext):
|
||||
def enterTanhFunc(self, ctx: MathExprParser.TanhFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#TanhFunc.
|
||||
def exitTanhFunc(self, ctx:MathExprParser.TanhFuncContext):
|
||||
def exitTanhFunc(self, ctx: MathExprParser.TanhFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AsinhFunc.
|
||||
def enterAsinhFunc(self, ctx:MathExprParser.AsinhFuncContext):
|
||||
def enterAsinhFunc(self, ctx: MathExprParser.AsinhFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AsinhFunc.
|
||||
def exitAsinhFunc(self, ctx:MathExprParser.AsinhFuncContext):
|
||||
def exitAsinhFunc(self, ctx: MathExprParser.AsinhFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AcoshFunc.
|
||||
def enterAcoshFunc(self, ctx:MathExprParser.AcoshFuncContext):
|
||||
def enterAcoshFunc(self, ctx: MathExprParser.AcoshFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AcoshFunc.
|
||||
def exitAcoshFunc(self, ctx:MathExprParser.AcoshFuncContext):
|
||||
def exitAcoshFunc(self, ctx: MathExprParser.AcoshFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AtanhFunc.
|
||||
def enterAtanhFunc(self, ctx:MathExprParser.AtanhFuncContext):
|
||||
def enterAtanhFunc(self, ctx: MathExprParser.AtanhFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AtanhFunc.
|
||||
def exitAtanhFunc(self, ctx:MathExprParser.AtanhFuncContext):
|
||||
def exitAtanhFunc(self, ctx: MathExprParser.AtanhFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#AbsFunc.
|
||||
def enterAbsFunc(self, ctx:MathExprParser.AbsFuncContext):
|
||||
def enterAbsFunc(self, ctx: MathExprParser.AbsFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#AbsFunc.
|
||||
def exitAbsFunc(self, ctx:MathExprParser.AbsFuncContext):
|
||||
def exitAbsFunc(self, ctx: MathExprParser.AbsFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SqrtFunc.
|
||||
def enterSqrtFunc(self, ctx:MathExprParser.SqrtFuncContext):
|
||||
def enterSqrtFunc(self, ctx: MathExprParser.SqrtFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SqrtFunc.
|
||||
def exitSqrtFunc(self, ctx:MathExprParser.SqrtFuncContext):
|
||||
def exitSqrtFunc(self, ctx: MathExprParser.SqrtFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#LnFunc.
|
||||
def enterLnFunc(self, ctx:MathExprParser.LnFuncContext):
|
||||
def enterLnFunc(self, ctx: MathExprParser.LnFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#LnFunc.
|
||||
def exitLnFunc(self, ctx:MathExprParser.LnFuncContext):
|
||||
def exitLnFunc(self, ctx: MathExprParser.LnFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#LogFunc.
|
||||
def enterLogFunc(self, ctx:MathExprParser.LogFuncContext):
|
||||
def enterLogFunc(self, ctx: MathExprParser.LogFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#LogFunc.
|
||||
def exitLogFunc(self, ctx:MathExprParser.LogFuncContext):
|
||||
def exitLogFunc(self, ctx: MathExprParser.LogFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ExpFunc.
|
||||
def enterExpFunc(self, ctx:MathExprParser.ExpFuncContext):
|
||||
def enterExpFunc(self, ctx: MathExprParser.ExpFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ExpFunc.
|
||||
def exitExpFunc(self, ctx:MathExprParser.ExpFuncContext):
|
||||
def exitExpFunc(self, ctx: MathExprParser.ExpFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#TNormFunc.
|
||||
def enterTNormFunc(self, ctx:MathExprParser.TNormFuncContext):
|
||||
def enterTNormFunc(self, ctx: MathExprParser.TNormFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#TNormFunc.
|
||||
def exitTNormFunc(self, ctx:MathExprParser.TNormFuncContext):
|
||||
def exitTNormFunc(self, ctx: MathExprParser.TNormFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SNormFunc.
|
||||
def enterSNormFunc(self, ctx:MathExprParser.SNormFuncContext):
|
||||
def enterSNormFunc(self, ctx: MathExprParser.SNormFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SNormFunc.
|
||||
def exitSNormFunc(self, ctx:MathExprParser.SNormFuncContext):
|
||||
def exitSNormFunc(self, ctx: MathExprParser.SNormFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#FloorFunc.
|
||||
def enterFloorFunc(self, ctx:MathExprParser.FloorFuncContext):
|
||||
def enterFloorFunc(self, ctx: MathExprParser.FloorFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#FloorFunc.
|
||||
def exitFloorFunc(self, ctx:MathExprParser.FloorFuncContext):
|
||||
def exitFloorFunc(self, ctx: MathExprParser.FloorFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#CeilFunc.
|
||||
def enterCeilFunc(self, ctx:MathExprParser.CeilFuncContext):
|
||||
def enterCeilFunc(self, ctx: MathExprParser.CeilFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#CeilFunc.
|
||||
def exitCeilFunc(self, ctx:MathExprParser.CeilFuncContext):
|
||||
def exitCeilFunc(self, ctx: MathExprParser.CeilFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#RoundFunc.
|
||||
def enterRoundFunc(self, ctx:MathExprParser.RoundFuncContext):
|
||||
def enterRoundFunc(self, ctx: MathExprParser.RoundFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#RoundFunc.
|
||||
def exitRoundFunc(self, ctx:MathExprParser.RoundFuncContext):
|
||||
def exitRoundFunc(self, ctx: MathExprParser.RoundFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#GammaFunc.
|
||||
def enterGammaFunc(self, ctx:MathExprParser.GammaFuncContext):
|
||||
def enterGammaFunc(self, ctx: MathExprParser.GammaFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#GammaFunc.
|
||||
def exitGammaFunc(self, ctx:MathExprParser.GammaFuncContext):
|
||||
def exitGammaFunc(self, ctx: MathExprParser.GammaFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#sigmoidFunc.
|
||||
def enterSigmoidFunc(self, ctx:MathExprParser.SigmoidFuncContext):
|
||||
def enterSigmoidFunc(self, ctx: MathExprParser.SigmoidFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#sigmoidFunc.
|
||||
def exitSigmoidFunc(self, ctx:MathExprParser.SigmoidFuncContext):
|
||||
def exitSigmoidFunc(self, ctx: MathExprParser.SigmoidFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#sfftFunc.
|
||||
def enterSfftFunc(self, ctx:MathExprParser.SfftFuncContext):
|
||||
def enterSfftFunc(self, ctx: MathExprParser.SfftFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#sfftFunc.
|
||||
def exitSfftFunc(self, ctx:MathExprParser.SfftFuncContext):
|
||||
def exitSfftFunc(self, ctx: MathExprParser.SfftFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#sifftFunc.
|
||||
def enterSifftFunc(self, ctx:MathExprParser.SifftFuncContext):
|
||||
def enterSifftFunc(self, ctx: MathExprParser.SifftFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#sifftFunc.
|
||||
def exitSifftFunc(self, ctx:MathExprParser.SifftFuncContext):
|
||||
def exitSifftFunc(self, ctx: MathExprParser.SifftFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#anglFunc.
|
||||
def enterAnglFunc(self, ctx:MathExprParser.AnglFuncContext):
|
||||
def enterAnglFunc(self, ctx: MathExprParser.AnglFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#anglFunc.
|
||||
def exitAnglFunc(self, ctx:MathExprParser.AnglFuncContext):
|
||||
def exitAnglFunc(self, ctx: MathExprParser.AnglFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#printFunc.
|
||||
def enterPrintFunc(self, ctx:MathExprParser.PrintFuncContext):
|
||||
def enterPrintFunc(self, ctx: MathExprParser.PrintFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#printFunc.
|
||||
def exitPrintFunc(self, ctx:MathExprParser.PrintFuncContext):
|
||||
def exitPrintFunc(self, ctx: MathExprParser.PrintFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#FractFunc.
|
||||
def enterFractFunc(self, ctx:MathExprParser.FractFuncContext):
|
||||
def enterFractFunc(self, ctx: MathExprParser.FractFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#FractFunc.
|
||||
def exitFractFunc(self, ctx:MathExprParser.FractFuncContext):
|
||||
def exitFractFunc(self, ctx: MathExprParser.FractFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ReluFunc.
|
||||
def enterReluFunc(self, ctx:MathExprParser.ReluFuncContext):
|
||||
def enterReluFunc(self, ctx: MathExprParser.ReluFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ReluFunc.
|
||||
def exitReluFunc(self, ctx:MathExprParser.ReluFuncContext):
|
||||
def exitReluFunc(self, ctx: MathExprParser.ReluFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SoftplusFunc.
|
||||
def enterSoftplusFunc(self, ctx:MathExprParser.SoftplusFuncContext):
|
||||
def enterSoftplusFunc(self, ctx: MathExprParser.SoftplusFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SoftplusFunc.
|
||||
def exitSoftplusFunc(self, ctx:MathExprParser.SoftplusFuncContext):
|
||||
def exitSoftplusFunc(self, ctx: MathExprParser.SoftplusFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#GeluFunc.
|
||||
def enterGeluFunc(self, ctx:MathExprParser.GeluFuncContext):
|
||||
def enterGeluFunc(self, ctx: MathExprParser.GeluFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#GeluFunc.
|
||||
def exitGeluFunc(self, ctx:MathExprParser.GeluFuncContext):
|
||||
def exitGeluFunc(self, ctx: MathExprParser.GeluFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SignFunc.
|
||||
def enterSignFunc(self, ctx:MathExprParser.SignFuncContext):
|
||||
def enterSignFunc(self, ctx: MathExprParser.SignFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SignFunc.
|
||||
def exitSignFunc(self, ctx:MathExprParser.SignFuncContext):
|
||||
def exitSignFunc(self, ctx: MathExprParser.SignFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#PrintShapeFunc.
|
||||
def enterPrintShapeFunc(self, ctx:MathExprParser.PrintShapeFuncContext):
|
||||
def enterPrintShapeFunc(self, ctx: MathExprParser.PrintShapeFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#PrintShapeFunc.
|
||||
def exitPrintShapeFunc(self, ctx:MathExprParser.PrintShapeFuncContext):
|
||||
def exitPrintShapeFunc(self, ctx: MathExprParser.PrintShapeFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#PowFunc.
|
||||
def enterPowFunc(self, ctx:MathExprParser.PowFuncContext):
|
||||
def enterPowFunc(self, ctx: MathExprParser.PowFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#PowFunc.
|
||||
def exitPowFunc(self, ctx:MathExprParser.PowFuncContext):
|
||||
def exitPowFunc(self, ctx: MathExprParser.PowFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#Atan2Func.
|
||||
def enterAtan2Func(self, ctx:MathExprParser.Atan2FuncContext):
|
||||
def enterAtan2Func(self, ctx: MathExprParser.Atan2FuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#Atan2Func.
|
||||
def exitAtan2Func(self, ctx:MathExprParser.Atan2FuncContext):
|
||||
def exitAtan2Func(self, ctx: MathExprParser.Atan2FuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#TMinFunc.
|
||||
def enterTMinFunc(self, ctx:MathExprParser.TMinFuncContext):
|
||||
def enterTMinFunc(self, ctx: MathExprParser.TMinFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#TMinFunc.
|
||||
def exitTMinFunc(self, ctx:MathExprParser.TMinFuncContext):
|
||||
def exitTMinFunc(self, ctx: MathExprParser.TMinFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#TMaxFunc.
|
||||
def enterTMaxFunc(self, ctx:MathExprParser.TMaxFuncContext):
|
||||
def enterTMaxFunc(self, ctx: MathExprParser.TMaxFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#TMaxFunc.
|
||||
def exitTMaxFunc(self, ctx:MathExprParser.TMaxFuncContext):
|
||||
def exitTMaxFunc(self, ctx: MathExprParser.TMaxFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#StepFunc.
|
||||
def enterStepFunc(self, ctx:MathExprParser.StepFuncContext):
|
||||
def enterStepFunc(self, ctx: MathExprParser.StepFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#StepFunc.
|
||||
def exitStepFunc(self, ctx:MathExprParser.StepFuncContext):
|
||||
def exitStepFunc(self, ctx: MathExprParser.StepFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ClampFunc.
|
||||
def enterClampFunc(self, ctx:MathExprParser.ClampFuncContext):
|
||||
def enterClampFunc(self, ctx: MathExprParser.ClampFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ClampFunc.
|
||||
def exitClampFunc(self, ctx:MathExprParser.ClampFuncContext):
|
||||
def exitClampFunc(self, ctx: MathExprParser.ClampFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#LerpFunc.
|
||||
def enterLerpFunc(self, ctx:MathExprParser.LerpFuncContext):
|
||||
def enterLerpFunc(self, ctx: MathExprParser.LerpFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#LerpFunc.
|
||||
def exitLerpFunc(self, ctx:MathExprParser.LerpFuncContext):
|
||||
def exitLerpFunc(self, ctx: MathExprParser.LerpFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SmoothstepFunc.
|
||||
def enterSmoothstepFunc(self, ctx:MathExprParser.SmoothstepFuncContext):
|
||||
def enterSmoothstepFunc(self, ctx: MathExprParser.SmoothstepFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SmoothstepFunc.
|
||||
def exitSmoothstepFunc(self, ctx:MathExprParser.SmoothstepFuncContext):
|
||||
def exitSmoothstepFunc(self, ctx: MathExprParser.SmoothstepFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#RangeFunc.
|
||||
def enterRangeFunc(self, ctx:MathExprParser.RangeFuncContext):
|
||||
def enterRangeFunc(self, ctx: MathExprParser.RangeFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#RangeFunc.
|
||||
def exitRangeFunc(self, ctx:MathExprParser.RangeFuncContext):
|
||||
def exitRangeFunc(self, ctx: MathExprParser.RangeFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SwapFunc.
|
||||
def enterSwapFunc(self, ctx:MathExprParser.SwapFuncContext):
|
||||
def enterSwapFunc(self, ctx: MathExprParser.SwapFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SwapFunc.
|
||||
def exitSwapFunc(self, ctx:MathExprParser.SwapFuncContext):
|
||||
def exitSwapFunc(self, ctx: MathExprParser.SwapFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SMinFunc.
|
||||
def enterSMinFunc(self, ctx:MathExprParser.SMinFuncContext):
|
||||
def enterSMinFunc(self, ctx: MathExprParser.SMinFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SMinFunc.
|
||||
def exitSMinFunc(self, ctx:MathExprParser.SMinFuncContext):
|
||||
def exitSMinFunc(self, ctx: MathExprParser.SMinFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#SMaxFunc.
|
||||
def enterSMaxFunc(self, ctx:MathExprParser.SMaxFuncContext):
|
||||
def enterSMaxFunc(self, ctx: MathExprParser.SMaxFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#SMaxFunc.
|
||||
def exitSMaxFunc(self, ctx:MathExprParser.SMaxFuncContext):
|
||||
def exitSMaxFunc(self, ctx: MathExprParser.SMaxFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#MapFunc.
|
||||
def enterMapFunc(self, ctx:MathExprParser.MapFuncContext):
|
||||
def enterMapFunc(self, ctx: MathExprParser.MapFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#MapFunc.
|
||||
def exitMapFunc(self, ctx:MathExprParser.MapFuncContext):
|
||||
def exitMapFunc(self, ctx: MathExprParser.MapFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ConvFunc.
|
||||
def enterConvFunc(self, ctx:MathExprParser.ConvFuncContext):
|
||||
def enterConvFunc(self, ctx: MathExprParser.ConvFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ConvFunc.
|
||||
def exitConvFunc(self, ctx:MathExprParser.ConvFuncContext):
|
||||
def exitConvFunc(self, ctx: MathExprParser.ConvFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#PermuteFunc.
|
||||
def enterPermuteFunc(self, ctx:MathExprParser.PermuteFuncContext):
|
||||
def enterPermuteFunc(self, ctx: MathExprParser.PermuteFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#PermuteFunc.
|
||||
def exitPermuteFunc(self, ctx:MathExprParser.PermuteFuncContext):
|
||||
def exitPermuteFunc(self, ctx: MathExprParser.PermuteFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
# Enter a parse tree produced by MathExprParser#ReshapeFunc.
|
||||
def enterReshapeFunc(self, ctx:MathExprParser.ReshapeFuncContext):
|
||||
def enterReshapeFunc(self, ctx: MathExprParser.ReshapeFuncContext):
|
||||
pass
|
||||
|
||||
# Exit a parse tree produced by MathExprParser#ReshapeFunc.
|
||||
def exitReshapeFunc(self, ctx:MathExprParser.ReshapeFuncContext):
|
||||
def exitReshapeFunc(self, ctx: MathExprParser.ReshapeFuncContext):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
del MathExprParser
|
||||
del MathExprParser
|
||||
|
||||
+5247
-1374
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,6 @@
|
||||
# Generated from ./MathExpr.g4 by ANTLR 4.13.2
|
||||
from antlr4 import *
|
||||
|
||||
if "." in __name__:
|
||||
from .MathExprParser import MathExprParser
|
||||
else:
|
||||
@@ -7,412 +8,331 @@ else:
|
||||
|
||||
# This class defines a complete generic visitor for a parse tree produced by MathExprParser.
|
||||
|
||||
|
||||
class MathExprVisitor(ParseTreeVisitor):
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#expr.
|
||||
def visitExpr(self, ctx:MathExprParser.ExprContext):
|
||||
def visitExpr(self, ctx: MathExprParser.ExprContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#LtExp.
|
||||
def visitLtExp(self, ctx:MathExprParser.LtExpContext):
|
||||
def visitLtExp(self, ctx: MathExprParser.LtExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#EqExp.
|
||||
def visitEqExp(self, ctx:MathExprParser.EqExpContext):
|
||||
def visitEqExp(self, ctx: MathExprParser.EqExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ToAdd.
|
||||
def visitToAdd(self, ctx:MathExprParser.ToAddContext):
|
||||
def visitToAdd(self, ctx: MathExprParser.ToAddContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#GeExp.
|
||||
def visitGeExp(self, ctx:MathExprParser.GeExpContext):
|
||||
def visitGeExp(self, ctx: MathExprParser.GeExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#LeExp.
|
||||
def visitLeExp(self, ctx:MathExprParser.LeExpContext):
|
||||
def visitLeExp(self, ctx: MathExprParser.LeExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#NeExp.
|
||||
def visitNeExp(self, ctx:MathExprParser.NeExpContext):
|
||||
def visitNeExp(self, ctx: MathExprParser.NeExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#GtExp.
|
||||
def visitGtExp(self, ctx:MathExprParser.GtExpContext):
|
||||
def visitGtExp(self, ctx: MathExprParser.GtExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AddExp.
|
||||
def visitAddExp(self, ctx:MathExprParser.AddExpContext):
|
||||
def visitAddExp(self, ctx: MathExprParser.AddExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ToMul.
|
||||
def visitToMul(self, ctx:MathExprParser.ToMulContext):
|
||||
def visitToMul(self, ctx: MathExprParser.ToMulContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SubExp.
|
||||
def visitSubExp(self, ctx:MathExprParser.SubExpContext):
|
||||
def visitSubExp(self, ctx: MathExprParser.SubExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#MulExp.
|
||||
def visitMulExp(self, ctx:MathExprParser.MulExpContext):
|
||||
def visitMulExp(self, ctx: MathExprParser.MulExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ModExp.
|
||||
def visitModExp(self, ctx:MathExprParser.ModExpContext):
|
||||
def visitModExp(self, ctx: MathExprParser.ModExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#DivExp.
|
||||
def visitDivExp(self, ctx:MathExprParser.DivExpContext):
|
||||
def visitDivExp(self, ctx: MathExprParser.DivExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ToPow.
|
||||
def visitToPow(self, ctx:MathExprParser.ToPowContext):
|
||||
def visitToPow(self, ctx: MathExprParser.ToPowContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#PowExp.
|
||||
def visitPowExp(self, ctx:MathExprParser.PowExpContext):
|
||||
def visitPowExp(self, ctx: MathExprParser.PowExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ToUnary.
|
||||
def visitToUnary(self, ctx:MathExprParser.ToUnaryContext):
|
||||
def visitToUnary(self, ctx: MathExprParser.ToUnaryContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#UnaryPlus.
|
||||
def visitUnaryPlus(self, ctx:MathExprParser.UnaryPlusContext):
|
||||
def visitUnaryPlus(self, ctx: MathExprParser.UnaryPlusContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#UnaryMinus.
|
||||
def visitUnaryMinus(self, ctx:MathExprParser.UnaryMinusContext):
|
||||
def visitUnaryMinus(self, ctx: MathExprParser.UnaryMinusContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ToAtom.
|
||||
def visitToAtom(self, ctx:MathExprParser.ToAtomContext):
|
||||
def visitToAtom(self, ctx: MathExprParser.ToAtomContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func1Exp.
|
||||
def visitFunc1Exp(self, ctx:MathExprParser.Func1ExpContext):
|
||||
def visitFunc1Exp(self, ctx: MathExprParser.Func1ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func2Exp.
|
||||
def visitFunc2Exp(self, ctx:MathExprParser.Func2ExpContext):
|
||||
def visitFunc2Exp(self, ctx: MathExprParser.Func2ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func3Exp.
|
||||
def visitFunc3Exp(self, ctx:MathExprParser.Func3ExpContext):
|
||||
def visitFunc3Exp(self, ctx: MathExprParser.Func3ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Func4Exp.
|
||||
def visitFunc4Exp(self, ctx:MathExprParser.Func4ExpContext):
|
||||
def visitFunc4Exp(self, ctx: MathExprParser.Func4ExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#FuncNExp.
|
||||
def visitFuncNExp(self, ctx:MathExprParser.FuncNExpContext):
|
||||
def visitFuncNExp(self, ctx: MathExprParser.FuncNExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#VariableExp.
|
||||
def visitVariableExp(self, ctx:MathExprParser.VariableExpContext):
|
||||
def visitVariableExp(self, ctx: MathExprParser.VariableExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#NumberExp.
|
||||
def visitNumberExp(self, ctx:MathExprParser.NumberExpContext):
|
||||
def visitNumberExp(self, ctx: MathExprParser.NumberExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ConstantExp.
|
||||
def visitConstantExp(self, ctx:MathExprParser.ConstantExpContext):
|
||||
def visitConstantExp(self, ctx: MathExprParser.ConstantExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ParenExp.
|
||||
def visitParenExp(self, ctx:MathExprParser.ParenExpContext):
|
||||
def visitParenExp(self, ctx: MathExprParser.ParenExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AbsExp.
|
||||
def visitAbsExp(self, ctx:MathExprParser.AbsExpContext):
|
||||
def visitAbsExp(self, ctx: MathExprParser.AbsExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ListExp.
|
||||
def visitListExp(self, ctx:MathExprParser.ListExpContext):
|
||||
def visitListExp(self, ctx: MathExprParser.ListExpContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SinFunc.
|
||||
def visitSinFunc(self, ctx:MathExprParser.SinFuncContext):
|
||||
def visitSinFunc(self, ctx: MathExprParser.SinFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#CosFunc.
|
||||
def visitCosFunc(self, ctx:MathExprParser.CosFuncContext):
|
||||
def visitCosFunc(self, ctx: MathExprParser.CosFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#TanFunc.
|
||||
def visitTanFunc(self, ctx:MathExprParser.TanFuncContext):
|
||||
def visitTanFunc(self, ctx: MathExprParser.TanFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AsinFunc.
|
||||
def visitAsinFunc(self, ctx:MathExprParser.AsinFuncContext):
|
||||
def visitAsinFunc(self, ctx: MathExprParser.AsinFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AcosFunc.
|
||||
def visitAcosFunc(self, ctx:MathExprParser.AcosFuncContext):
|
||||
def visitAcosFunc(self, ctx: MathExprParser.AcosFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AtanFunc.
|
||||
def visitAtanFunc(self, ctx:MathExprParser.AtanFuncContext):
|
||||
def visitAtanFunc(self, ctx: MathExprParser.AtanFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SinhFunc.
|
||||
def visitSinhFunc(self, ctx:MathExprParser.SinhFuncContext):
|
||||
def visitSinhFunc(self, ctx: MathExprParser.SinhFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#CoshFunc.
|
||||
def visitCoshFunc(self, ctx:MathExprParser.CoshFuncContext):
|
||||
def visitCoshFunc(self, ctx: MathExprParser.CoshFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#TanhFunc.
|
||||
def visitTanhFunc(self, ctx:MathExprParser.TanhFuncContext):
|
||||
def visitTanhFunc(self, ctx: MathExprParser.TanhFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AsinhFunc.
|
||||
def visitAsinhFunc(self, ctx:MathExprParser.AsinhFuncContext):
|
||||
def visitAsinhFunc(self, ctx: MathExprParser.AsinhFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AcoshFunc.
|
||||
def visitAcoshFunc(self, ctx:MathExprParser.AcoshFuncContext):
|
||||
def visitAcoshFunc(self, ctx: MathExprParser.AcoshFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AtanhFunc.
|
||||
def visitAtanhFunc(self, ctx:MathExprParser.AtanhFuncContext):
|
||||
def visitAtanhFunc(self, ctx: MathExprParser.AtanhFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#AbsFunc.
|
||||
def visitAbsFunc(self, ctx:MathExprParser.AbsFuncContext):
|
||||
def visitAbsFunc(self, ctx: MathExprParser.AbsFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SqrtFunc.
|
||||
def visitSqrtFunc(self, ctx:MathExprParser.SqrtFuncContext):
|
||||
def visitSqrtFunc(self, ctx: MathExprParser.SqrtFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#LnFunc.
|
||||
def visitLnFunc(self, ctx:MathExprParser.LnFuncContext):
|
||||
def visitLnFunc(self, ctx: MathExprParser.LnFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#LogFunc.
|
||||
def visitLogFunc(self, ctx:MathExprParser.LogFuncContext):
|
||||
def visitLogFunc(self, ctx: MathExprParser.LogFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ExpFunc.
|
||||
def visitExpFunc(self, ctx:MathExprParser.ExpFuncContext):
|
||||
def visitExpFunc(self, ctx: MathExprParser.ExpFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#TNormFunc.
|
||||
def visitTNormFunc(self, ctx:MathExprParser.TNormFuncContext):
|
||||
def visitTNormFunc(self, ctx: MathExprParser.TNormFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SNormFunc.
|
||||
def visitSNormFunc(self, ctx:MathExprParser.SNormFuncContext):
|
||||
def visitSNormFunc(self, ctx: MathExprParser.SNormFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#FloorFunc.
|
||||
def visitFloorFunc(self, ctx:MathExprParser.FloorFuncContext):
|
||||
def visitFloorFunc(self, ctx: MathExprParser.FloorFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#CeilFunc.
|
||||
def visitCeilFunc(self, ctx:MathExprParser.CeilFuncContext):
|
||||
def visitCeilFunc(self, ctx: MathExprParser.CeilFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#RoundFunc.
|
||||
def visitRoundFunc(self, ctx:MathExprParser.RoundFuncContext):
|
||||
def visitRoundFunc(self, ctx: MathExprParser.RoundFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#GammaFunc.
|
||||
def visitGammaFunc(self, ctx:MathExprParser.GammaFuncContext):
|
||||
def visitGammaFunc(self, ctx: MathExprParser.GammaFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#sigmoidFunc.
|
||||
def visitSigmoidFunc(self, ctx:MathExprParser.SigmoidFuncContext):
|
||||
def visitSigmoidFunc(self, ctx: MathExprParser.SigmoidFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#sfftFunc.
|
||||
def visitSfftFunc(self, ctx:MathExprParser.SfftFuncContext):
|
||||
def visitSfftFunc(self, ctx: MathExprParser.SfftFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#sifftFunc.
|
||||
def visitSifftFunc(self, ctx:MathExprParser.SifftFuncContext):
|
||||
def visitSifftFunc(self, ctx: MathExprParser.SifftFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#anglFunc.
|
||||
def visitAnglFunc(self, ctx:MathExprParser.AnglFuncContext):
|
||||
def visitAnglFunc(self, ctx: MathExprParser.AnglFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#printFunc.
|
||||
def visitPrintFunc(self, ctx:MathExprParser.PrintFuncContext):
|
||||
def visitPrintFunc(self, ctx: MathExprParser.PrintFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#FractFunc.
|
||||
def visitFractFunc(self, ctx:MathExprParser.FractFuncContext):
|
||||
def visitFractFunc(self, ctx: MathExprParser.FractFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ReluFunc.
|
||||
def visitReluFunc(self, ctx:MathExprParser.ReluFuncContext):
|
||||
def visitReluFunc(self, ctx: MathExprParser.ReluFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SoftplusFunc.
|
||||
def visitSoftplusFunc(self, ctx:MathExprParser.SoftplusFuncContext):
|
||||
def visitSoftplusFunc(self, ctx: MathExprParser.SoftplusFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#GeluFunc.
|
||||
def visitGeluFunc(self, ctx:MathExprParser.GeluFuncContext):
|
||||
def visitGeluFunc(self, ctx: MathExprParser.GeluFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SignFunc.
|
||||
def visitSignFunc(self, ctx:MathExprParser.SignFuncContext):
|
||||
def visitSignFunc(self, ctx: MathExprParser.SignFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#PrintShapeFunc.
|
||||
def visitPrintShapeFunc(self, ctx:MathExprParser.PrintShapeFuncContext):
|
||||
def visitPrintShapeFunc(self, ctx: MathExprParser.PrintShapeFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#PowFunc.
|
||||
def visitPowFunc(self, ctx:MathExprParser.PowFuncContext):
|
||||
def visitPowFunc(self, ctx: MathExprParser.PowFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#Atan2Func.
|
||||
def visitAtan2Func(self, ctx:MathExprParser.Atan2FuncContext):
|
||||
def visitAtan2Func(self, ctx: MathExprParser.Atan2FuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#TMinFunc.
|
||||
def visitTMinFunc(self, ctx:MathExprParser.TMinFuncContext):
|
||||
def visitTMinFunc(self, ctx: MathExprParser.TMinFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#TMaxFunc.
|
||||
def visitTMaxFunc(self, ctx:MathExprParser.TMaxFuncContext):
|
||||
def visitTMaxFunc(self, ctx: MathExprParser.TMaxFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#StepFunc.
|
||||
def visitStepFunc(self, ctx:MathExprParser.StepFuncContext):
|
||||
def visitStepFunc(self, ctx: MathExprParser.StepFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ClampFunc.
|
||||
def visitClampFunc(self, ctx:MathExprParser.ClampFuncContext):
|
||||
def visitClampFunc(self, ctx: MathExprParser.ClampFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#LerpFunc.
|
||||
def visitLerpFunc(self, ctx:MathExprParser.LerpFuncContext):
|
||||
def visitLerpFunc(self, ctx: MathExprParser.LerpFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SmoothstepFunc.
|
||||
def visitSmoothstepFunc(self, ctx:MathExprParser.SmoothstepFuncContext):
|
||||
def visitSmoothstepFunc(self, ctx: MathExprParser.SmoothstepFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#RangeFunc.
|
||||
def visitRangeFunc(self, ctx:MathExprParser.RangeFuncContext):
|
||||
def visitRangeFunc(self, ctx: MathExprParser.RangeFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SwapFunc.
|
||||
def visitSwapFunc(self, ctx:MathExprParser.SwapFuncContext):
|
||||
def visitSwapFunc(self, ctx: MathExprParser.SwapFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SMinFunc.
|
||||
def visitSMinFunc(self, ctx:MathExprParser.SMinFuncContext):
|
||||
def visitSMinFunc(self, ctx: MathExprParser.SMinFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#SMaxFunc.
|
||||
def visitSMaxFunc(self, ctx:MathExprParser.SMaxFuncContext):
|
||||
def visitSMaxFunc(self, ctx: MathExprParser.SMaxFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#MapFunc.
|
||||
def visitMapFunc(self, ctx:MathExprParser.MapFuncContext):
|
||||
def visitMapFunc(self, ctx: MathExprParser.MapFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ConvFunc.
|
||||
def visitConvFunc(self, ctx:MathExprParser.ConvFuncContext):
|
||||
def visitConvFunc(self, ctx: MathExprParser.ConvFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#PermuteFunc.
|
||||
def visitPermuteFunc(self, ctx:MathExprParser.PermuteFuncContext):
|
||||
def visitPermuteFunc(self, ctx: MathExprParser.PermuteFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
# Visit a parse tree produced by MathExprParser#ReshapeFunc.
|
||||
def visitReshapeFunc(self, ctx:MathExprParser.ReshapeFuncContext):
|
||||
def visitReshapeFunc(self, ctx: MathExprParser.ReshapeFuncContext):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
|
||||
|
||||
del MathExprParser
|
||||
del MathExprParser
|
||||
|
||||
@@ -4,22 +4,28 @@ import torch.nn.functional as F
|
||||
from .MathExprVisitor import MathExprVisitor
|
||||
from ..helper_functions import generate_dim_variables, getIndexTensorAlongDim
|
||||
|
||||
|
||||
class UnifiedMathVisitor(MathExprVisitor):
|
||||
def __init__(self, variables, shape=None, device=None):
|
||||
self.variables = variables
|
||||
self.spatial_variables = variables.copy()
|
||||
self.shape = shape if shape is not None else ()
|
||||
if device is None:
|
||||
self.device = next((v.device for v in variables.values() if isinstance(v, torch.Tensor)), torch.device("cpu"))
|
||||
self.device = next((v.device for v in variables.values() if isinstance(v, torch.Tensor)), torch.device("cpu"))
|
||||
else:
|
||||
self.device = device
|
||||
self.device = device
|
||||
|
||||
def _is_tensor(self, val): return isinstance(val, torch.Tensor)
|
||||
def _is_list(self, val): return isinstance(val, (list, tuple))
|
||||
def _is_tensor(self, val):
|
||||
return isinstance(val, torch.Tensor)
|
||||
|
||||
def _is_list(self, val):
|
||||
return isinstance(val, (list, tuple))
|
||||
|
||||
def _promote_to_tensor(self, val):
|
||||
if self._is_tensor(val): return val
|
||||
if self._is_list(val): return torch.tensor(val, device=self.device)
|
||||
if self._is_tensor(val):
|
||||
return val
|
||||
if self._is_list(val):
|
||||
return torch.tensor(val, device=self.device)
|
||||
return torch.tensor(val, device=self.device)
|
||||
|
||||
def _bin_op(self, a, b, torch_op, scalar_op):
|
||||
@@ -27,17 +33,20 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
Generic binary operation handler.
|
||||
"""
|
||||
# one of them is a list and one is tensor
|
||||
if self._is_tensor(a) and self._is_list(b): return torch.cat([self._bin_op(a, x, torch_op, scalar_op) for x in b], dim=0)
|
||||
if self._is_list(a) and self._is_tensor(b): return torch.cat([self._bin_op(x, b, torch_op, scalar_op) for x in a], dim=0)
|
||||
if self._is_tensor(a) and self._is_list(b):
|
||||
return torch.cat([self._bin_op(a, x, torch_op, scalar_op) for x in b], dim=0)
|
||||
if self._is_list(a) and self._is_tensor(b):
|
||||
return torch.cat([self._bin_op(x, b, torch_op, scalar_op) for x in a], dim=0)
|
||||
|
||||
if self._is_list(a) and not self._is_tensor(b):
|
||||
if self._is_list(b):
|
||||
if len(a) != len(b): raise ValueError("List length mismatch")
|
||||
if len(a) != len(b):
|
||||
raise ValueError("List length mismatch")
|
||||
return [self._bin_op(x, y, torch_op, scalar_op) for x, y in zip(a, b)]
|
||||
return [self._bin_op(x, b, torch_op, scalar_op) for x in a]
|
||||
|
||||
if not self._is_tensor(a) and self._is_list(b):
|
||||
return [self._bin_op(a, x, torch_op, scalar_op) for x in b]
|
||||
return [self._bin_op(a, x, torch_op, scalar_op) for x in b]
|
||||
|
||||
if self._is_tensor(a) or self._is_tensor(b):
|
||||
if torch_op:
|
||||
@@ -59,14 +68,16 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
|
||||
def visitNumberExp(self, ctx):
|
||||
val_str = ctx.getText()
|
||||
if '.' in val_str or 'e' in val_str:
|
||||
if "." in val_str or "e" in val_str:
|
||||
return float(val_str)
|
||||
return int(val_str)
|
||||
|
||||
def visitConstantExp(self, ctx):
|
||||
name = ctx.getText().lower()
|
||||
if name == "pi": return math.pi
|
||||
if name == "e": return math.e
|
||||
if name == "pi":
|
||||
return math.pi
|
||||
if name == "e":
|
||||
return math.e
|
||||
raise ValueError(f"Unknown constant: {name}")
|
||||
|
||||
def visitVariableExp(self, ctx):
|
||||
@@ -89,39 +100,43 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
|
||||
# Binary Ops
|
||||
def visitAddExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()),
|
||||
torch.add, lambda a,b: a+b)
|
||||
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()), torch.add, lambda a, b: a + b)
|
||||
|
||||
def visitSubExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()),
|
||||
torch.sub, lambda a,b: a-b)
|
||||
return self._bin_op(self.visit(ctx.addExpr()), self.visit(ctx.mulExpr()), torch.sub, lambda a, b: a - b)
|
||||
|
||||
def visitMulExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()),
|
||||
torch.mul, lambda a,b: a*b)
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.mul, lambda a, b: a * b)
|
||||
|
||||
def visitDivExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()),
|
||||
torch.div, lambda a,b: a/b)
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.div, lambda a, b: a / b)
|
||||
|
||||
def visitModExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()),
|
||||
torch.fmod, lambda a,b: a%b)
|
||||
return self._bin_op(self.visit(ctx.mulExpr()), self.visit(ctx.powExpr()), torch.fmod, lambda a, b: a % b)
|
||||
|
||||
def visitPowExp(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.unaryExpr()), self.visit(ctx.powExpr()),
|
||||
torch.pow, math.pow)
|
||||
|
||||
return self._bin_op(self.visit(ctx.unaryExpr()), self.visit(ctx.powExpr()), torch.pow, math.pow)
|
||||
|
||||
def _bool_op(self, a, b, torch_op, scalar_op):
|
||||
return self._bin_op(a, b, torch_op, scalar_op)
|
||||
|
||||
def visitNeExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ne, lambda a,b: float(a!=b))
|
||||
def visitEqExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.eq, lambda a,b: float(a==b))
|
||||
def visitGtExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.gt, lambda a,b: float(a>b))
|
||||
def visitLtExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.lt, lambda a,b: float(a<b))
|
||||
def visitGeExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ge, lambda a,b: float(a>=b))
|
||||
def visitLeExp(self, ctx): return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.le, lambda a,b: float(a<=b))
|
||||
def visitNeExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ne, lambda a, b: float(a != b))
|
||||
|
||||
def visitEqExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.eq, lambda a, b: float(a == b))
|
||||
|
||||
def visitGtExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.gt, lambda a, b: float(a > b))
|
||||
|
||||
def visitLtExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.lt, lambda a, b: float(a < b))
|
||||
|
||||
def visitGeExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.ge, lambda a, b: float(a >= b))
|
||||
|
||||
def visitLeExp(self, ctx):
|
||||
return self._bool_op(self.visit(ctx.compExpr()), self.visit(ctx.addExpr()), torch.le, lambda a, b: float(a <= b))
|
||||
|
||||
# Functions
|
||||
def _func_dispatch(self, arg, torch_fn, scalar_fn):
|
||||
@@ -131,55 +146,107 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return torch_fn(arg)
|
||||
return scalar_fn(arg)
|
||||
|
||||
def visitSinFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.sin, math.sin)
|
||||
def visitCosFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.cos, math.cos)
|
||||
def visitTanFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.tan, math.tan)
|
||||
def visitAsinFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.asin, math.asin)
|
||||
def visitAcosFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.acos, math.acos)
|
||||
def visitAtanFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.atan, math.atan)
|
||||
def visitSinhFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.sinh, math.sinh)
|
||||
def visitCoshFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.cosh, math.cosh)
|
||||
def visitTanhFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.tanh, math.tanh)
|
||||
def visitAsinhFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.asinh, math.asinh)
|
||||
def visitAcoshFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.acosh, math.acosh)
|
||||
def visitAtanhFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.atanh, math.atanh)
|
||||
def visitSinFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.sin, math.sin)
|
||||
|
||||
def visitCosFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.cos, math.cos)
|
||||
|
||||
def visitTanFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.tan, math.tan)
|
||||
|
||||
def visitAsinFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.asin, math.asin)
|
||||
|
||||
def visitAcosFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.acos, math.acos)
|
||||
|
||||
def visitAtanFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.atan, math.atan)
|
||||
|
||||
def visitSinhFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.sinh, math.sinh)
|
||||
|
||||
def visitCoshFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.cosh, math.cosh)
|
||||
|
||||
def visitTanhFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.tanh, math.tanh)
|
||||
|
||||
def visitAsinhFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.asinh, math.asinh)
|
||||
|
||||
def visitAcoshFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.acosh, math.acosh)
|
||||
|
||||
def visitAtanhFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.atanh, math.atanh)
|
||||
|
||||
def visitAbsFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.abs, abs)
|
||||
|
||||
def visitAbsFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.abs, abs)
|
||||
def visitAbsExp(self, ctx):
|
||||
val = self.visit(ctx.expr())
|
||||
if self._is_list(val): return torch.linalg.norm(self._promote_to_tensor(val))
|
||||
if self._is_tensor(val): return torch.linalg.norm(val)
|
||||
if self._is_list(val):
|
||||
return torch.linalg.norm(self._promote_to_tensor(val))
|
||||
if self._is_tensor(val):
|
||||
return torch.linalg.norm(val)
|
||||
return abs(val)
|
||||
|
||||
def visitSqrtFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.sqrt, math.sqrt)
|
||||
def visitLnFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.log, math.log)
|
||||
def visitLogFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.log10, math.log10)
|
||||
def visitExpFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.exp, math.exp)
|
||||
def visitSqrtFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.sqrt, math.sqrt)
|
||||
|
||||
def visitFloorFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.floor, math.floor)
|
||||
def visitCeilFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.ceil, math.ceil)
|
||||
def visitRoundFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.round, round)
|
||||
def visitLnFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.log, math.log)
|
||||
|
||||
def visitSignFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.sign, lambda x: math.copysign(1.0, x))
|
||||
def visitLogFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.log10, math.log10)
|
||||
|
||||
def visitExpFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.exp, math.exp)
|
||||
|
||||
def visitFloorFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.floor, math.floor)
|
||||
|
||||
def visitCeilFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.ceil, math.ceil)
|
||||
|
||||
def visitRoundFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.round, round)
|
||||
|
||||
def visitSignFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.sign, lambda x: math.copysign(1.0, x))
|
||||
|
||||
def visitFractFunc(self, ctx):
|
||||
val = self.visit(ctx.expr())
|
||||
if self._is_list(val): return [x - math.floor(x) for x in val]
|
||||
if self._is_tensor(val): return val - torch.floor(val)
|
||||
if self._is_list(val):
|
||||
return [x - math.floor(x) for x in val]
|
||||
if self._is_tensor(val):
|
||||
return val - torch.floor(val)
|
||||
return val - math.floor(val)
|
||||
|
||||
def visitGammaFunc(self, ctx):
|
||||
torch_gamma = getattr(torch.special, 'gamma', None)
|
||||
torch_gamma = getattr(torch.special, "gamma", None)
|
||||
if torch_gamma is None:
|
||||
torch_gamma = lambda x: torch.exp(torch.lgamma(x))
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch_gamma, math.gamma)
|
||||
|
||||
def visitSigmoidFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.sigmoid, lambda x: 1.0 / (1.0 + math.exp(-x)))
|
||||
def visitReluFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.relu, lambda x: max(0.0, x))
|
||||
def visitSoftplusFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), F.softplus, lambda x: math.log(1.0 + math.exp(x)))
|
||||
def visitGeluFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), F.gelu, lambda x: 0.5 * x * (1 + math.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * math.pow(x, 3)))))
|
||||
def visitSigmoidFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.sigmoid, lambda x: 1.0 / (1.0 + math.exp(-x)))
|
||||
|
||||
def visitAnglFunc(self, ctx): return self._func_dispatch(self.visit(ctx.expr()), torch.angle, lambda x: math.atan2(0, x) if x < 0 else 0)
|
||||
def visitReluFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.relu, lambda x: max(0.0, x))
|
||||
|
||||
def visitSoftplusFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), F.softplus, lambda x: math.log(1.0 + math.exp(x)))
|
||||
|
||||
def visitGeluFunc(self, ctx):
|
||||
return self._func_dispatch(
|
||||
self.visit(ctx.expr()), F.gelu, lambda x: 0.5 * x * (1 + math.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * math.pow(x, 3))))
|
||||
)
|
||||
|
||||
def visitAnglFunc(self, ctx):
|
||||
return self._func_dispatch(self.visit(ctx.expr()), torch.angle, lambda x: math.atan2(0, x) if x < 0 else 0)
|
||||
|
||||
def visitPrintFunc(self, ctx):
|
||||
val = self.visit(ctx.expr())
|
||||
@@ -188,25 +255,37 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
|
||||
def visitTNormFunc(self, ctx):
|
||||
val = self.visit(ctx.expr())
|
||||
if self._is_tensor(val): return F.normalize(val, p=2, dim=-1)
|
||||
if self._is_tensor(val):
|
||||
return F.normalize(val, p=2, dim=-1)
|
||||
return 1.0 if val != 0 else 0.0
|
||||
|
||||
def visitSNormFunc(self, ctx):
|
||||
val = self.visit(ctx.expr())
|
||||
if self._is_tensor(val): return torch.linalg.norm(val)
|
||||
if self._is_tensor(val):
|
||||
return torch.linalg.norm(val)
|
||||
return abs(val)
|
||||
|
||||
# Two-argument functions
|
||||
def visitPowFunc(self, ctx): return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.pow, math.pow)
|
||||
def visitAtan2Func(self, ctx): return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.atan2, math.atan2)
|
||||
def visitTMinFunc(self, ctx): return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.minimum, min)
|
||||
def visitTMaxFunc(self, ctx): return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.maximum, max)
|
||||
def visitPowFunc(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.pow, math.pow)
|
||||
|
||||
def visitAtan2Func(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.atan2, math.atan2)
|
||||
|
||||
def visitTMinFunc(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.minimum, min)
|
||||
|
||||
def visitTMaxFunc(self, ctx):
|
||||
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), torch.maximum, max)
|
||||
|
||||
def visitStepFunc(self, ctx):
|
||||
# step(x, edge) = 1 if x >= edge else 0
|
||||
return self._bin_op(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)),
|
||||
lambda x, edge: torch.where(x >= edge, 1.0, 0.0),
|
||||
lambda x, edge: 1.0 if x >= edge else 0.0)
|
||||
return self._bin_op(
|
||||
self.visit(ctx.expr(0)),
|
||||
self.visit(ctx.expr(1)),
|
||||
lambda x, edge: torch.where(x >= edge, 1.0, 0.0),
|
||||
lambda x, edge: 1.0 if x >= edge else 0.0,
|
||||
)
|
||||
|
||||
# Three-argument functions
|
||||
def visitClampFunc(self, ctx):
|
||||
@@ -217,7 +296,7 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
if any(self._is_tensor(x) for x in [val, min_v, max_v]):
|
||||
return torch.clamp(self._promote_to_tensor(val), self._promote_to_tensor(min_v), self._promote_to_tensor(max_v))
|
||||
if self._is_list(val):
|
||||
return [max(min(x, max_v), min_v) for x in val] #TODO: what if min_v or max_v is list
|
||||
return [max(min(x, max_v), min_v) for x in val] # TODO: what if min_v or max_v is list
|
||||
return max(min(val, max_v), min_v)
|
||||
|
||||
def visitLerpFunc(self, ctx):
|
||||
@@ -249,11 +328,20 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return list(torch.arange(self.visit(ctx.expr(0)), self.visit(ctx.expr(1)), self.visit(ctx.expr(2))))
|
||||
|
||||
# Helpers for visiting generic exprs
|
||||
def visitFunc1Exp(self, ctx): return self.visitChildren(ctx)
|
||||
def visitFunc2Exp(self, ctx): return self.visitChildren(ctx)
|
||||
def visitFuncNExp(self, ctx): return self.visitChildren(ctx)
|
||||
def visitAtomExp(self, ctx): return self.visitChildren(ctx)
|
||||
def visitExpr(self, ctx): return self.visitChildren(ctx)
|
||||
def visitFunc1Exp(self, ctx):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
def visitFunc2Exp(self, ctx):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
def visitFuncNExp(self, ctx):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
def visitAtomExp(self, ctx):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
def visitExpr(self, ctx):
|
||||
return self.visitChildren(ctx)
|
||||
|
||||
def visitSMinFunc(self, ctx):
|
||||
vals = [self.visit(e) for e in ctx.expr()]
|
||||
@@ -262,15 +350,18 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return min(vals)
|
||||
|
||||
promoted = [self._promote_to_tensor(x) for x in vals]
|
||||
if len(promoted) == 1: return torch.min(promoted[0])
|
||||
if len(promoted) == 1:
|
||||
return torch.min(promoted[0])
|
||||
return torch.min(torch.stack(torch.broadcast_tensors(*promoted)))
|
||||
|
||||
def visitSMaxFunc(self, ctx):
|
||||
args = [self.visit(e) for e in ctx.expr()]
|
||||
if len(args) == 1:
|
||||
# Check if list or scalar
|
||||
if not self._is_tensor(args[0]) and not self._is_list(args[0]): return args[0]
|
||||
if self._is_list(args[0]): return max(args[0]) # max of list
|
||||
if not self._is_tensor(args[0]) and not self._is_list(args[0]):
|
||||
return args[0]
|
||||
if self._is_list(args[0]):
|
||||
return max(args[0]) # max of list
|
||||
return torch.max(args[0]) # Global max of single tensor
|
||||
|
||||
# Multiple args
|
||||
@@ -278,7 +369,8 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return max(args)
|
||||
|
||||
promoted = [self._promote_to_tensor(x) for x in args]
|
||||
if len(promoted) == 1: return torch.max(promoted[0])
|
||||
if len(promoted) == 1:
|
||||
return torch.max(promoted[0])
|
||||
return torch.max(torch.stack(torch.broadcast_tensors(*promoted)))
|
||||
|
||||
def _fold_nd(self, tsr, spatial_dims):
|
||||
@@ -294,10 +386,10 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
new_batch = 1
|
||||
for i in range(fold_count + 1):
|
||||
new_batch *= tsr.shape[i]
|
||||
tsr = tsr.reshape(new_batch, *tsr.shape[fold_count+1:])
|
||||
tsr = tsr.reshape(new_batch, *tsr.shape[fold_count + 1 :])
|
||||
folded = True
|
||||
else:
|
||||
folded = (added_dims > 0)
|
||||
folded = added_dims > 0
|
||||
return tsr, original_shape, added_dims, folded
|
||||
|
||||
def _unfold_nd(self, tsr, original_shape, added_dims, folded):
|
||||
@@ -315,20 +407,22 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
tsr = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
||||
dims = self.visit(ctx.expr(1))
|
||||
if isinstance(dims, torch.Tensor):
|
||||
dims = dims.flatten().long().tolist()
|
||||
dims = dims.flatten().long().tolist()
|
||||
return tsr.permute(*dims)
|
||||
|
||||
def visitReshapeFunc(self, ctx):
|
||||
tsr = self._promote_to_tensor(self.visit(ctx.expr(0)))
|
||||
new_shape = self.visit(ctx.expr(1))
|
||||
if isinstance(new_shape, torch.Tensor):
|
||||
new_shape = new_shape.flatten().long().tolist()
|
||||
new_shape = new_shape.flatten().long().tolist()
|
||||
return tsr.reshape(*new_shape)
|
||||
|
||||
def visitPrintShapeFunc(self, ctx):
|
||||
tsr = self.visit(ctx.expr())
|
||||
if hasattr(tsr, 'shape'): print(tsr.shape)
|
||||
else: print(f"Scalar/List: {tsr}")
|
||||
if hasattr(tsr, "shape"):
|
||||
print(tsr.shape)
|
||||
else:
|
||||
print(f"Scalar/List: {tsr}")
|
||||
return tsr
|
||||
|
||||
def visitSfftFunc(self, ctx):
|
||||
@@ -346,10 +440,10 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
self.variables = self.variables.copy()
|
||||
device = self.device
|
||||
|
||||
shape_to_use = self.shape if self.shape else (1,1,1,1)
|
||||
shape_to_use = self.shape if self.shape else (1, 1, 1, 1)
|
||||
|
||||
ndim = len(shape_to_use)
|
||||
dim_names = ['x', 'y', 'z', 'w', 'v', 'u']
|
||||
dim_names = ["x", "y", "z", "w", "v", "u"]
|
||||
|
||||
k_components = []
|
||||
for i in range(ndim):
|
||||
@@ -361,23 +455,23 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
values = values.view(*view_shape).expand(*shape_to_use)
|
||||
|
||||
if i < len(dim_names):
|
||||
var_name = f'K{dim_names[i]}'
|
||||
var_name = f"K{dim_names[i]}"
|
||||
self.variables[var_name] = values
|
||||
self.variables[f'F{dim_names[i]}'] = float(size_d)
|
||||
self.variables[f"F{dim_names[i]}"] = float(size_d)
|
||||
|
||||
self.variables[f'K_dim{dim_idx}'] = values
|
||||
self.variables[f'F_dim{dim_idx}'] = float(size_d)
|
||||
self.variables[f"K_dim{dim_idx}"] = values
|
||||
self.variables[f"F_dim{dim_idx}"] = float(size_d)
|
||||
k_components.append(values)
|
||||
|
||||
k_sq_sum = torch.zeros(shape_to_use, device=device)
|
||||
for k_val in k_components:
|
||||
k_sq_sum = k_sq_sum + k_val ** 2
|
||||
k_sq_sum = k_sq_sum + k_val**2
|
||||
|
||||
self.variables['K'] = torch.sqrt(k_sq_sum)
|
||||
self.variables['frequency'] = self.variables['K']
|
||||
self.variables["K"] = torch.sqrt(k_sq_sum)
|
||||
self.variables["frequency"] = self.variables["K"]
|
||||
|
||||
if 'Kx' in self.variables:
|
||||
self.variables['frequency_count'] = self.variables.get('Fx', 1.0)
|
||||
if "Kx" in self.variables:
|
||||
self.variables["frequency_count"] = self.variables.get("Fx", 1.0)
|
||||
|
||||
self.variables = self.variables | generate_dim_variables(k_sq_sum)
|
||||
|
||||
@@ -398,9 +492,12 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
i = int(idx1_t.flatten()[0].item()) if isinstance(idx1_t, torch.Tensor) else int(idx1_t)
|
||||
j = int(idx2_t.flatten()[0].item()) if isinstance(idx2_t, torch.Tensor) else int(idx2_t)
|
||||
|
||||
while dim < 0: dim += tsr.ndim
|
||||
while i < 0: i += tsr.shape[dim]
|
||||
while j < 0: j += tsr.shape[dim]
|
||||
while dim < 0:
|
||||
dim += tsr.ndim
|
||||
while i < 0:
|
||||
i += tsr.shape[dim]
|
||||
while j < 0:
|
||||
j += tsr.shape[dim]
|
||||
|
||||
indices = torch.arange(tsr.shape[dim], device=tsr.device)
|
||||
val_i = indices[i].clone()
|
||||
@@ -410,7 +507,8 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
return torch.index_select(tsr, dim, indices)
|
||||
|
||||
def _normalize_coord(self, coord, size):
|
||||
if size > 1: return (coord / (size - 1)) * 2.0 - 1.0
|
||||
if size > 1:
|
||||
return (coord / (size - 1)) * 2.0 - 1.0
|
||||
return torch.zeros_like(coord)
|
||||
|
||||
def visitMapFunc(self, ctx):
|
||||
@@ -418,14 +516,17 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
coords = [self._promote_to_tensor(self.visit(ctx.expr(i))) for i in range(1, len(ctx.expr()))]
|
||||
num_coords = len(coords)
|
||||
|
||||
if num_coords == 0: return tensor
|
||||
if num_coords > 3: raise ValueError("map() supports max 3 mapping functions.")
|
||||
if num_coords == 0:
|
||||
return tensor
|
||||
if num_coords > 3:
|
||||
raise ValueError("map() supports max 3 mapping functions.")
|
||||
|
||||
spatial_in_shape = tensor.shape[-num_coords:]
|
||||
leading_shape = tensor.shape[:-num_coords]
|
||||
|
||||
batch_size = 1
|
||||
for s in leading_shape: batch_size *= s
|
||||
for s in leading_shape:
|
||||
batch_size *= s
|
||||
|
||||
input_view = tensor.reshape(batch_size, 1, *spatial_in_shape)
|
||||
norm_coords_list = []
|
||||
@@ -441,14 +542,15 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
grid_view = grid.reshape(batch_size, *grid_spatial_shape[-(num_coords):], num_coords)
|
||||
except RuntimeError:
|
||||
grid_view = grid.expand(batch_size, *([-1] * len(grid_spatial_shape)), -1)
|
||||
grid_view = grid_view.reshape(batch_size, *grid_view.shape[-(num_coords+1):-1], num_coords)
|
||||
grid_view = grid_view.reshape(batch_size, *grid_view.shape[-(num_coords + 1) : -1], num_coords)
|
||||
|
||||
if num_coords == 1:
|
||||
input_final = input_view.reshape(batch_size, 1, 1, -1)
|
||||
gv = grid_view
|
||||
while gv.ndim < 3: gv = gv.unsqueeze(1)
|
||||
while gv.ndim < 3:
|
||||
gv = gv.unsqueeze(1)
|
||||
y_zeros = torch.zeros_like(gv[..., :1])
|
||||
grid_final = torch.cat([gv, y_zeros], dim=-1).unsqueeze(1) # [B, 1, 1, 2]
|
||||
grid_final = torch.cat([gv, y_zeros], dim=-1).unsqueeze(1) # [B, 1, 1, 2]
|
||||
output = F.grid_sample(input_final, grid_final, align_corners=True)
|
||||
elif num_coords == 2:
|
||||
grid_final = grid_view.reshape(batch_size, *grid_view.shape[-3:-1], 2)
|
||||
@@ -476,18 +578,18 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
p_high = p_total - p_low
|
||||
pads.extend([p_low, p_high])
|
||||
|
||||
padded_input = torch.nn.functional.pad(conv_input, tuple(pads), mode='constant', value=0)
|
||||
padded_input = torch.nn.functional.pad(conv_input, tuple(pads), mode="constant", value=0)
|
||||
|
||||
actual_kernel_sizes = kernel_sizes[::-1]
|
||||
|
||||
if kernel_val.numel() == 1:
|
||||
kernel_val = kernel_val.expand(tuple(actual_kernel_sizes))
|
||||
elif kernel_val.ndim != spatial_dims_count:
|
||||
kernel_val = kernel_val.reshape(tuple(actual_kernel_sizes))
|
||||
kernel_val = kernel_val.reshape(tuple(actual_kernel_sizes))
|
||||
|
||||
final_kernel = kernel_val.unsqueeze(0).unsqueeze(0)
|
||||
final_kernel = final_kernel.to(conv_input.dtype)
|
||||
final_kernel = final_kernel.repeat(in_channels, 1, *([1]*spatial_dims_count))
|
||||
final_kernel = final_kernel.repeat(in_channels, 1, *([1] * spatial_dims_count))
|
||||
|
||||
conv_fn = F.conv1d if spatial_dims_count == 1 else (F.conv2d if spatial_dims_count == 2 else F.conv3d)
|
||||
|
||||
@@ -499,33 +601,35 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
tensor = self._promote_to_tensor(input_raw)
|
||||
input_ndim = tensor.ndim
|
||||
num_args = len(ctx.expr())
|
||||
if num_args < 3: raise ValueError("conv() requires at least 3 arguments")
|
||||
if num_args < 3:
|
||||
raise ValueError("conv() requires at least 3 arguments")
|
||||
|
||||
kernel_arg_idx = num_args - 1
|
||||
|
||||
spatial_dims_count = num_args - 2
|
||||
if spatial_dims_count not in [1, 2, 3]:
|
||||
raise ValueError(f"conv() supports 1D, 2D, or 3D. Found {spatial_dims_count}")
|
||||
raise ValueError(f"conv() supports 1D, 2D, or 3D. Found {spatial_dims_count}")
|
||||
|
||||
kernel_sizes = []
|
||||
for i in range(1, 1 + spatial_dims_count):
|
||||
val = self.visit(ctx.expr(i))
|
||||
if isinstance(val, torch.Tensor): val = int(val.flatten()[0].item())
|
||||
if isinstance(val, torch.Tensor):
|
||||
val = int(val.flatten()[0].item())
|
||||
kernel_sizes.append(val)
|
||||
|
||||
# Prepare context for kernel
|
||||
coords = [torch.arange(s, device=self.device).float() - (s//2) for s in kernel_sizes]
|
||||
grid = torch.meshgrid(*coords, indexing='ij')
|
||||
coords = [torch.arange(s, device=self.device).float() - (s // 2) for s in kernel_sizes]
|
||||
grid = torch.meshgrid(*coords, indexing="ij")
|
||||
|
||||
dim_names = ['x','y','z']
|
||||
dim_names = ["x", "y", "z"]
|
||||
old_vars = self.variables
|
||||
self.variables = self.variables.copy()
|
||||
for i in range(spatial_dims_count):
|
||||
self.variables[f'k{dim_names[i].lower()}'] = grid[i]
|
||||
self.variables[f'k{dim_names[i].upper()}'] = grid[i]
|
||||
self.variables[f"k{dim_names[i].lower()}"] = grid[i]
|
||||
self.variables[f"k{dim_names[i].upper()}"] = grid[i]
|
||||
|
||||
kernel_var_names = ['kW', 'kH', 'kD']
|
||||
kernel_var_full = ['kernel_width', 'kernel_height', 'kernel_depth']
|
||||
kernel_var_names = ["kW", "kH", "kD"]
|
||||
kernel_var_full = ["kernel_width", "kernel_height", "kernel_depth"]
|
||||
for i in range(spatial_dims_count):
|
||||
size_val = float(kernel_sizes[i])
|
||||
self.variables[kernel_var_names[i]] = size_val
|
||||
@@ -543,27 +647,29 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
is_channels_first = False
|
||||
if tensor.ndim == spatial_dims_count + 2:
|
||||
c_front = tensor.shape[1]
|
||||
c_back = tensor.shape[-1]
|
||||
c_back = tensor.shape[-1]
|
||||
if c_front <= 4 and c_front < c_back:
|
||||
is_channels_first = True
|
||||
elif c_back <= 4 and c_back < c_front:
|
||||
is_channels_first = False
|
||||
else:
|
||||
is_channels_first = (c_back > 4)
|
||||
is_channels_first = c_back > 4
|
||||
|
||||
if is_channels_first:
|
||||
if spatial_dims_count == 1 and tensor.ndim == 3:
|
||||
tensor = tensor.permute(0, 2, 1)
|
||||
tensor = tensor.permute(0, 2, 1)
|
||||
elif spatial_dims_count == 2 and tensor.ndim == 4:
|
||||
tensor = tensor.permute(0, 2, 3, 1)
|
||||
tensor = tensor.permute(0, 2, 3, 1)
|
||||
elif spatial_dims_count == 3:
|
||||
if tensor.ndim == 4: tensor = tensor.unsqueeze(-1)
|
||||
elif tensor.ndim == 5: tensor = tensor.permute(0, 2, 3, 4, 1)
|
||||
if tensor.ndim == 4:
|
||||
tensor = tensor.unsqueeze(-1)
|
||||
elif tensor.ndim == 5:
|
||||
tensor = tensor.permute(0, 2, 3, 4, 1)
|
||||
|
||||
if tensor.ndim == spatial_dims_count + 1:
|
||||
tensor = tensor.unsqueeze(-1) # Add C=1
|
||||
tensor = tensor.unsqueeze(-1) # Add C=1
|
||||
elif tensor.ndim == spatial_dims_count:
|
||||
tensor = tensor.unsqueeze(0).unsqueeze(-1) # Add B=1, C=1
|
||||
tensor = tensor.unsqueeze(0).unsqueeze(-1) # Add B=1, C=1
|
||||
|
||||
in_channels = tensor.size(-1)
|
||||
batch_end_idx = tensor.ndim - 1 - spatial_dims_count
|
||||
@@ -572,7 +678,8 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
channels_shape = (tensor.shape[-1],)
|
||||
|
||||
total_batch = 1
|
||||
for s in batch_shape: total_batch *= s
|
||||
for s in batch_shape:
|
||||
total_batch *= s
|
||||
flat_input = tensor.reshape(total_batch, *spatial_shape, in_channels)
|
||||
|
||||
permute_order = [0, spatial_dims_count + 1] + list(range(1, spatial_dims_count + 1))
|
||||
@@ -580,22 +687,21 @@ class UnifiedMathVisitor(MathExprVisitor):
|
||||
|
||||
result = self._apply_conv_internal(conv_input, kernel_val, kernel_sizes, spatial_dims_count)
|
||||
|
||||
result_permute = [0] + list(range(2, 2+spatial_dims_count)) + [1]
|
||||
result_permute = [0] + list(range(2, 2 + spatial_dims_count)) + [1]
|
||||
out_flat = result.permute(*result_permute)
|
||||
|
||||
final_shape = batch_shape + spatial_shape + channels_shape
|
||||
out = out_flat.reshape(final_shape)
|
||||
|
||||
if is_channels_first:
|
||||
if spatial_dims_count == 1 and out.ndim == 3:
|
||||
out = out.permute(0, 2, 1)
|
||||
elif spatial_dims_count == 2 and out.ndim == 4:
|
||||
out = out.permute(0, 3, 1, 2)
|
||||
elif spatial_dims_count == 3:
|
||||
if out.ndim == 5 and out.shape[-1] == 1 and input_ndim == 4:
|
||||
out = out.squeeze(-1)
|
||||
elif out.ndim == 5:
|
||||
out = out.permute(0, 4, 1, 2, 3)
|
||||
if spatial_dims_count == 1 and out.ndim == 3:
|
||||
out = out.permute(0, 2, 1)
|
||||
elif spatial_dims_count == 2 and out.ndim == 4:
|
||||
out = out.permute(0, 3, 1, 2)
|
||||
elif spatial_dims_count == 3:
|
||||
if out.ndim == 5 and out.shape[-1] == 1 and input_ndim == 4:
|
||||
out = out.squeeze(-1)
|
||||
elif out.ndim == 5:
|
||||
out = out.permute(0, 4, 1, 2, 3)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@@ -1,13 +1,15 @@
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
windows = {'bartlet':torch.bartlett_window, 'blackman':torch.blackman_window, 'hamming':torch.hamming_window,'hann':torch.hann_window}
|
||||
windows = {"bartlet": torch.bartlett_window, "blackman": torch.blackman_window, "hamming": torch.hamming_window, "hann": torch.hann_window}
|
||||
|
||||
|
||||
class SpectrogramToAudio(io.ComfyNode):
|
||||
"""
|
||||
Converts an Image spectrogram back to Audio.
|
||||
Red is real part and blue is imaginary. Green is ignored.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -21,7 +23,6 @@ class SpectrogramToAudio(io.ComfyNode):
|
||||
io.Int.Input(id="window_length", default=1024, min=16, tooltip="Window length in samples"),
|
||||
io.Int.Input(id="hop_length", default=256, min=1, tooltip="Stride of the window (hop length) in samples"),
|
||||
io.Combo.Input(id="window_type", default="hann", options=list(windows.keys()), tooltip="Type of window function to apply"),
|
||||
|
||||
],
|
||||
outputs=[
|
||||
io.Audio.Output(id="audio", tooltip="Output audio"),
|
||||
@@ -29,19 +30,24 @@ class SpectrogramToAudio(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, image, channel_count, sample_rate, window_length, hop_length,window_type):
|
||||
def execute(cls, image, channel_count, sample_rate, window_length, hop_length, window_type):
|
||||
B, H, W, _ = image.shape
|
||||
bucket_count = H // channel_count
|
||||
n_fft = (bucket_count - 1) * 2
|
||||
real = image[..., 0].reshape(B * channel_count, bucket_count, W)
|
||||
imag = image[..., 2].reshape(B * channel_count, bucket_count, W)
|
||||
stft_complex = torch.complex(real, imag)
|
||||
window = windows[window_type](window_length,device = image.device)
|
||||
window = windows[window_type](window_length, device=image.device)
|
||||
waveform = torch.istft(
|
||||
stft_complex, n_fft=n_fft, hop_length=hop_length, win_length=window_length,
|
||||
window=window, center=True, normalized=False, onesided=True
|
||||
stft_complex,
|
||||
n_fft=n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=window_length,
|
||||
window=window,
|
||||
center=True,
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
)
|
||||
|
||||
waveform = waveform.reshape(B, channel_count, -1)
|
||||
return ({"waveform": waveform, "sample_rate": sample_rate},)
|
||||
|
||||
|
||||
@@ -4,10 +4,12 @@ import copy
|
||||
from .modelLikeCommon import calculate_patches
|
||||
from .helper_functions import comonLazy
|
||||
|
||||
|
||||
class VAEMathNode(io.ComfyNode):
|
||||
"""
|
||||
This node enables the use of math expressions on VAE weights.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
@@ -16,13 +18,13 @@ class VAEMathNode(io.ComfyNode):
|
||||
category="More math",
|
||||
inputs=[
|
||||
io.Vae.Input(id="a", tooltip="Main VAE (base)"),
|
||||
io.Vae.Input(id="b", optional=True,lazy=True, tooltip="Optional 2nd VAE"),
|
||||
io.Vae.Input(id="c", optional=True,lazy=True, tooltip="Optional 3rd VAE"),
|
||||
io.Vae.Input(id="d", optional=True,lazy=True, tooltip="Optional 4th VAE"),
|
||||
io.Float.Input(id="w", default=0.0,lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0,lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="y", default=0.0,lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="z", default=0.0,lazy=True, optional=True, force_input=True),
|
||||
io.Vae.Input(id="b", optional=True, lazy=True, tooltip="Optional 2nd VAE"),
|
||||
io.Vae.Input(id="c", optional=True, lazy=True, tooltip="Optional 3rd VAE"),
|
||||
io.Vae.Input(id="d", optional=True, lazy=True, tooltip="Optional 4th VAE"),
|
||||
io.Float.Input(id="w", default=0.0, lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="x", default=0.0, lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="y", default=0.0, lazy=True, optional=True, force_input=True),
|
||||
io.Float.Input(id="z", default=0.0, lazy=True, optional=True, force_input=True),
|
||||
io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
|
||||
],
|
||||
outputs=[
|
||||
@@ -31,9 +33,11 @@ class VAEMathNode(io.ComfyNode):
|
||||
)
|
||||
|
||||
tooltip = cleandoc(__doc__)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
|
||||
return comonLazy(Model, a, b, c, d,w,x,y,z)
|
||||
def check_lazy_status(cls, Model, a, b=[], c=[], d=[], w=0, x=0, y=0, z=0):
|
||||
return comonLazy(Model, a, b, c, d, w, x, y, z)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
|
||||
patcher_a = a.patcher
|
||||
|
||||
+60
-32
@@ -11,16 +11,17 @@ from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, ev
|
||||
|
||||
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
|
||||
"""
|
||||
@@ -53,7 +54,7 @@ class VideoMathNode(MathNodeBase):
|
||||
@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)
|
||||
@@ -65,45 +66,72 @@ class VideoMathNode(MathNodeBase):
|
||||
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],
|
||||
"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_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],
|
||||
"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
|
||||
))
|
||||
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,)
|
||||
|
||||
@@ -7,18 +7,20 @@ from .Parser.MathExprLexer import MathExprLexer
|
||||
from .Parser.MathExprParser import MathExprParser
|
||||
|
||||
|
||||
|
||||
class ThrowingErrorListener(ErrorListener):
|
||||
"""Error listener that raises ValueError on syntax errors."""
|
||||
|
||||
def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e):
|
||||
raise ValueError(f"Syntax error in expression at line {line}, col {column}: {msg}")
|
||||
|
||||
def as_tensor(value,shape):
|
||||
if isinstance(value,torch.Tensor):
|
||||
|
||||
def as_tensor(value, shape):
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value
|
||||
if isinstance(value,(float,int)):
|
||||
if isinstance(value, (float, int)):
|
||||
value = (value,)
|
||||
return torch.broadcast_to(torch.Tensor(value),shape)
|
||||
return torch.broadcast_to(torch.Tensor(value), shape)
|
||||
|
||||
|
||||
def parse_expr(expr: str):
|
||||
"""Parse a math expression and return the parse tree."""
|
||||
@@ -33,6 +35,7 @@ def parse_expr(expr: str):
|
||||
def eval_tensor_expr(expr: str, variables: dict, shape: tuple, device=None):
|
||||
"""Parse and evaluate a tensor math expression."""
|
||||
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
|
||||
tree = parse_expr(expr)
|
||||
visitor = UnifiedMathVisitor(variables, shape, device=device)
|
||||
return visitor.visit(tree)
|
||||
@@ -41,6 +44,7 @@ def eval_tensor_expr(expr: str, variables: dict, shape: tuple, device=None):
|
||||
def eval_tensor_expr_with_tree(tree, variables: dict, shape: tuple, device=None):
|
||||
"""Evaluate a pre-parsed expression tree with UnifiedMathVisitor."""
|
||||
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
|
||||
visitor = UnifiedMathVisitor(variables, shape, device=device)
|
||||
return visitor.visit(tree)
|
||||
|
||||
@@ -48,14 +52,17 @@ def eval_tensor_expr_with_tree(tree, variables: dict, shape: tuple, device=None)
|
||||
def eval_float_expr(expr: str, variables: dict):
|
||||
"""Parse and evaluate a float math expression."""
|
||||
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
|
||||
tree = parse_expr(expr)
|
||||
# Float eval context often has no shape. Pass None/Empty.
|
||||
visitor = UnifiedMathVisitor(variables, shape=None)
|
||||
return visitor.visit(tree)
|
||||
|
||||
|
||||
def eval_float_expr_with_tree(tree, variables: dict):
|
||||
"""Evaluate a pre-parsed expression tree with UnifiedMathVisitor."""
|
||||
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
|
||||
visitor = UnifiedMathVisitor(variables, shape=None)
|
||||
return visitor.visit(tree)
|
||||
|
||||
@@ -72,7 +79,7 @@ def getIndexTensorAlongDim(tensor, dim):
|
||||
|
||||
def comonLazy(expr, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
"""Determine which lazy inputs are needed based on expression variables."""
|
||||
variables = {'a': a, 'b': b, 'c': c, 'd': d, 'w': w, 'x': x, 'y': y, 'z': z}
|
||||
variables = {"a": a, "b": b, "c": c, "d": d, "w": w, "x": x, "y": y, "z": z}
|
||||
need_eval = []
|
||||
input_stream = InputStream(expr)
|
||||
lexer = MathExprLexer(input_stream)
|
||||
@@ -83,14 +90,16 @@ def comonLazy(expr, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
need_eval.append(token.text)
|
||||
return need_eval
|
||||
|
||||
|
||||
def generate_dim_variables(tensor: torch.Tensor):
|
||||
"""Generate index and size tensors for each dimension of the input tensor."""
|
||||
variables = {}
|
||||
for dim, size in enumerate(tensor.shape):
|
||||
variables[f'D{dim}'] = getIndexTensorAlongDim(tensor, dim)
|
||||
variables[f'S{dim}'] = torch.full(tensor.shape, fill_value=size, dtype=torch.float32, device=tensor.device)
|
||||
variables[f"D{dim}"] = getIndexTensorAlongDim(tensor, dim)
|
||||
variables[f"S{dim}"] = torch.full(tensor.shape, fill_value=size, dtype=torch.float32, device=tensor.device)
|
||||
return variables
|
||||
|
||||
|
||||
def make_zero_like(ref):
|
||||
"""
|
||||
Create a zero-initialized version of the reference object, maintaining its structure.
|
||||
@@ -109,31 +118,23 @@ def make_zero_like(ref):
|
||||
# Ensure it's a tensor-like structure
|
||||
if torch.is_tensor(ref_tensor):
|
||||
ref_pooled = ref[0][1].get("pooled_output")
|
||||
return [[
|
||||
torch.zeros_like(ref_tensor),
|
||||
{"pooled_output": torch.zeros_like(ref_pooled) if ref_pooled is not None else None}
|
||||
]]
|
||||
return [[torch.zeros_like(ref_tensor), {"pooled_output": torch.zeros_like(ref_pooled) if ref_pooled is not None else None}]]
|
||||
|
||||
# Audio or Latent: dict
|
||||
if isinstance(ref, dict):
|
||||
if 'waveform' in ref: # Audio
|
||||
return {
|
||||
'waveform': torch.zeros_like(ref['waveform']),
|
||||
'sample_rate': ref['sample_rate']
|
||||
}
|
||||
if 'samples' in ref: # Latent
|
||||
return {
|
||||
'samples': torch.zeros_like(ref['samples'])
|
||||
}
|
||||
if "waveform" in ref: # Audio
|
||||
return {"waveform": torch.zeros_like(ref["waveform"]), "sample_rate": ref["sample_rate"]}
|
||||
if "samples" in ref: # Latent
|
||||
return {"samples": torch.zeros_like(ref["samples"])}
|
||||
|
||||
# VideoComponents or other objects with images/audio attributes
|
||||
if hasattr(ref, 'images') and hasattr(ref, 'audio'):
|
||||
if hasattr(ref, "images") and hasattr(ref, "audio"):
|
||||
# Dynamically create same type (e.g. VideoComponents)
|
||||
return type(ref)(
|
||||
images=torch.zeros_like(ref.images),
|
||||
audio={'waveform': torch.zeros_like(ref.audio['waveform']), 'sample_rate': ref.audio['sample_rate']},
|
||||
audio={"waveform": torch.zeros_like(ref.audio["waveform"]), "sample_rate": ref.audio["sample_rate"]},
|
||||
frame_rate=ref.frame_rate,
|
||||
metadata=None
|
||||
metadata=None,
|
||||
)
|
||||
|
||||
return None
|
||||
@@ -151,4 +152,4 @@ def freq_to_time(element: torch.Tensor) -> torch.Tensor:
|
||||
if element.ndim < 2:
|
||||
raise ValueError("IFFT requires at least 2 dimensions (Batch, Channel)")
|
||||
dims = tuple(range(2, element.ndim))
|
||||
return torch.fft.ifftn(element, dim=dims).real
|
||||
return torch.fft.ifftn(element, dim=dims).real
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
import comfy.utils
|
||||
|
||||
from .helper_functions import getIndexTensorAlongDim, parse_expr, eval_tensor_expr_with_tree,as_tensor
|
||||
from .helper_functions import getIndexTensorAlongDim, parse_expr, eval_tensor_expr_with_tree, as_tensor
|
||||
|
||||
|
||||
def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
||||
@@ -31,19 +31,27 @@ def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0
|
||||
|
||||
# Build variables
|
||||
variables = {
|
||||
'a': tens_a, 'b': tens_b, 'c': tens_c, 'd': tens_d,
|
||||
'w': w, 'x': x, 'y': y, 'z': z,
|
||||
'L': i, 'layer': i,
|
||||
'LC': layer_count, 'layer_count': layer_count,
|
||||
"a": tens_a,
|
||||
"b": tens_b,
|
||||
"c": tens_c,
|
||||
"d": tens_d,
|
||||
"w": w,
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z": z,
|
||||
"L": i,
|
||||
"layer": i,
|
||||
"LC": layer_count,
|
||||
"layer_count": layer_count,
|
||||
}
|
||||
|
||||
# Add dimension index tensors
|
||||
for dim_idx in range(tens_a.ndim):
|
||||
idx_tensor = getIndexTensorAlongDim(tens_a, dim_idx)
|
||||
variables[f'D{dim_idx}'] = idx_tensor
|
||||
variables[f'dim_{dim_idx}'] = idx_tensor
|
||||
variables[f"D{dim_idx}"] = idx_tensor
|
||||
variables[f"dim_{dim_idx}"] = idx_tensor
|
||||
|
||||
result_tensor = as_tensor(eval_tensor_expr_with_tree(tree, variables, tens_a.shape),tens_a.shape)
|
||||
result_tensor = as_tensor(eval_tensor_expr_with_tree(tree, variables, tens_a.shape), tens_a.shape)
|
||||
|
||||
# Calculate patch (diff from original)
|
||||
diff = result_tensor - tens_a
|
||||
@@ -54,4 +62,4 @@ def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
return patches
|
||||
return patches
|
||||
|
||||
+38
-31
@@ -13,37 +13,40 @@ from .AudioToSpectrogramNode import AudioToSpectrogram
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
class IntToFloatNode(io.ComfyNode):
|
||||
"""
|
||||
Converts int to float.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
"""
|
||||
"""
|
||||
return io.Schema(
|
||||
node_id="mrmth_IntToFloat",
|
||||
display_name="Int -> Float",
|
||||
category="More math",
|
||||
inputs=[
|
||||
io.Int.Input(id="value", default=0),
|
||||
],
|
||||
outputs=[
|
||||
io.Float.Output(),
|
||||
],
|
||||
)
|
||||
""" """
|
||||
return io.Schema(
|
||||
node_id="mrmth_IntToFloat",
|
||||
display_name="Int -> Float",
|
||||
category="More math",
|
||||
inputs=[
|
||||
io.Int.Input(id="value", default=0),
|
||||
],
|
||||
outputs=[
|
||||
io.Float.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, value):
|
||||
return (float(value),)
|
||||
|
||||
|
||||
class FloatToIntNode(io.ComfyNode):
|
||||
"""
|
||||
Converts float to int.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
"""
|
||||
"""
|
||||
""" """
|
||||
return io.Schema(
|
||||
node_id="mrmth_FloatToInt",
|
||||
category="More math",
|
||||
@@ -55,6 +58,7 @@ class FloatToIntNode(io.ComfyNode):
|
||||
io.Int.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, value):
|
||||
return (int(value),)
|
||||
@@ -63,23 +67,26 @@ class FloatToIntNode(io.ComfyNode):
|
||||
class MoreMathExtension(ComfyExtension):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
ConditioningMathNode,
|
||||
ModelMathNode,
|
||||
CLIPMathNode,
|
||||
VAEMathNode,
|
||||
LatentMathNode,
|
||||
ImageMathNode,
|
||||
FloatMathNode,
|
||||
NoiseMathNode,
|
||||
IntToFloatNode,
|
||||
FloatToIntNode,
|
||||
AudioMathNode,
|
||||
VideoMathNode,
|
||||
AudioToSpectrogram,
|
||||
SpectrogramToAudio
|
||||
]
|
||||
return [
|
||||
ConditioningMathNode,
|
||||
ModelMathNode,
|
||||
CLIPMathNode,
|
||||
VAEMathNode,
|
||||
LatentMathNode,
|
||||
ImageMathNode,
|
||||
FloatMathNode,
|
||||
NoiseMathNode,
|
||||
IntToFloatNode,
|
||||
FloatToIntNode,
|
||||
AudioMathNode,
|
||||
VideoMathNode,
|
||||
AudioToSpectrogram,
|
||||
SpectrogramToAudio,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> MoreMathExtension:
|
||||
return MoreMathExtension()
|
||||
|
||||
+1
-1
@@ -15,6 +15,7 @@ for p in [_project_root, _comfy_root]:
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def session_setup():
|
||||
"""
|
||||
@@ -23,4 +24,3 @@ def session_setup():
|
||||
ensures any further session setup can be added without changing tests.
|
||||
"""
|
||||
yield
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
|
||||
from more_math.LatentMathNode import LatentMathNode
|
||||
|
||||
|
||||
def test_conv_1d():
|
||||
"""
|
||||
Test 1D convolution.
|
||||
@@ -26,17 +27,17 @@ def test_conv_1d():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 10, 4)
|
||||
a_val = torch.randn(*shape)
|
||||
|
||||
|
||||
# conv(a, 3, 1.0) -> implies kernel of ones, size 3
|
||||
# Result should correspond to 1D conv
|
||||
try:
|
||||
# LatentMathNode expects latent dicts usually
|
||||
input_dict = {"samples": a_val}
|
||||
res = node.execute("conv(a, 3, 1.0)", a=input_dict)
|
||||
|
||||
|
||||
# LatentMathNode returns list of dicts
|
||||
res_tensor = res[0]["samples"]
|
||||
|
||||
|
||||
print(f"1D Conv Result Shape: {res_tensor.shape}")
|
||||
# Expect (1, 10, 4)
|
||||
assert res_tensor.shape == shape
|
||||
@@ -44,6 +45,7 @@ def test_conv_1d():
|
||||
print(f"1D Conv Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_3d():
|
||||
"""
|
||||
Test 3D convolution.
|
||||
@@ -54,7 +56,7 @@ def test_conv_3d():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 5, 32, 32, 4)
|
||||
a_val = torch.randn(*shape)
|
||||
|
||||
|
||||
try:
|
||||
input_dict = {"samples": a_val}
|
||||
res = node.execute("conv(a, 3, 3, 3, 1.0)", a=input_dict)
|
||||
@@ -65,6 +67,7 @@ def test_conv_3d():
|
||||
print(f"3D Conv Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_arbitrary_batch():
|
||||
"""
|
||||
Test generic tensor with extra batch dims.
|
||||
@@ -74,7 +77,7 @@ def test_conv_arbitrary_batch():
|
||||
node = LatentMathNode()
|
||||
shape = (2, 2, 16, 16, 4)
|
||||
a_val = torch.randn(*shape)
|
||||
|
||||
|
||||
try:
|
||||
# conv(a, 3, 3, 1.0) -> 2D conv on (16,16)
|
||||
input_dict = {"samples": a_val}
|
||||
@@ -86,6 +89,7 @@ def test_conv_arbitrary_batch():
|
||||
print(f"Arbitrary Batch Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_list_kernel():
|
||||
"""
|
||||
Test conv with list kernel (Regression test for float64 mismatch).
|
||||
@@ -93,15 +97,15 @@ def test_conv_list_kernel():
|
||||
"""
|
||||
print("\n--- Testing List Kernel Conv ---")
|
||||
node = LatentMathNode()
|
||||
shape = (1, 5, 10, 10, 4) # [B, D, H, W, C]
|
||||
shape = (1, 5, 10, 10, 4) # [B, D, H, W, C]
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
# 3x3x3 kernel = 27 elements
|
||||
# Using the user's example kernel
|
||||
kernel_list = [1,1,1,1,0,1,1,1,1, 0,0,0,0,1,0,0,0,0, 1,1,1,1,0,1,1,1,1]
|
||||
kernel_list = [1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1]
|
||||
kernel_str = str(kernel_list)
|
||||
expr = f"conv(a, 3, 3, 3, {kernel_str})/8"
|
||||
|
||||
|
||||
try:
|
||||
input_dict = {"samples": a_val}
|
||||
res = node.execute(expr, a=input_dict)
|
||||
@@ -113,6 +117,7 @@ def test_conv_list_kernel():
|
||||
print(f"List Kernel Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_audio():
|
||||
"""
|
||||
Test 1D conv on Audio [B, C, L].
|
||||
@@ -122,9 +127,9 @@ def test_conv_audio():
|
||||
"""
|
||||
print("\n--- Testing Audio Conv [B, C, L] ---")
|
||||
node = LatentMathNode()
|
||||
shape = (1, 2, 100) # [B, C, L] (L >> C)
|
||||
shape = (1, 2, 100) # [B, C, L] (L >> C)
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
# conv(a, 3, 1.0) on last dim (L)
|
||||
# Expected: result shape same as input
|
||||
try:
|
||||
@@ -132,19 +137,20 @@ def test_conv_audio():
|
||||
res = node.execute("conv(a, 3, 1.0)", a=input_dict)
|
||||
res_tensor = res[0]["samples"]
|
||||
print(f"Audio Result Shape: {res_tensor.shape}")
|
||||
|
||||
|
||||
if res_tensor.shape != shape:
|
||||
print(f"Likely interpreted as Channels Last [B, L, C] where C is small? No.")
|
||||
# If interpreted as Channels last [..., C].
|
||||
# [1, 2, 100]. Spatial=[2]. Channel=100.
|
||||
# Output [1, 2, 100] (but confusing channels).
|
||||
pass
|
||||
|
||||
print(f"Likely interpreted as Channels Last [B, L, C] where C is small? No.")
|
||||
# If interpreted as Channels last [..., C].
|
||||
# [1, 2, 100]. Spatial=[2]. Channel=100.
|
||||
# Output [1, 2, 100] (but confusing channels).
|
||||
pass
|
||||
|
||||
assert res_tensor.shape == shape
|
||||
except Exception as e:
|
||||
print(f"Audio Conv Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_deep_latent():
|
||||
"""
|
||||
Test 3D conv on Deep Latent [B, 32, H, W] (User request).
|
||||
@@ -156,7 +162,7 @@ def test_conv_deep_latent():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 32, 16, 16)
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
# conv(a, 3, 3, 3, 1.0)
|
||||
# 3D kernels need D,H,W.
|
||||
# D=32 (Channel). H=16. W=16.
|
||||
@@ -166,7 +172,7 @@ def test_conv_deep_latent():
|
||||
res_tensor = res[0]["samples"]
|
||||
print(f"Deep Latent Result Shape: {res_tensor.shape}")
|
||||
assert res_tensor.shape == shape
|
||||
|
||||
|
||||
# Identity check (ensure D neighbors engaged)
|
||||
# Using simple kernel, center only vs ones.
|
||||
# But this test just checks shape and execution path.
|
||||
@@ -174,6 +180,7 @@ def test_conv_deep_latent():
|
||||
print(f"Deep Latent Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_padding():
|
||||
"""
|
||||
Test padding consistency, especially for even kernels.
|
||||
@@ -184,7 +191,7 @@ def test_conv_padding():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 10, 10, 1)
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
# conv(a, 4, 4, 1.0)
|
||||
# If padding is symmetric 2, result is 11x11.
|
||||
# If padding is symmetric 1, result is 9x9.
|
||||
@@ -199,6 +206,7 @@ def test_conv_padding():
|
||||
print(f"Padding Test Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_complex_padding():
|
||||
"""
|
||||
Test asymmetric padding with mixed odd/even kernel sizes.
|
||||
@@ -209,7 +217,7 @@ def test_conv_complex_padding():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 10, 10, 1)
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
try:
|
||||
input_dict = {"samples": a_val}
|
||||
res = node.execute("conv(a, 3, 4, 1.0)", a=input_dict)
|
||||
@@ -220,6 +228,7 @@ def test_conv_complex_padding():
|
||||
print(f"Complex Padding Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_conv_3d_asymmetric():
|
||||
"""
|
||||
Test 3D conv with asymmetric spatial dims.
|
||||
@@ -229,7 +238,7 @@ def test_conv_3d_asymmetric():
|
||||
node = LatentMathNode()
|
||||
shape = (1, 5, 10, 20, 1)
|
||||
a_val = torch.randn(*shape).float()
|
||||
|
||||
|
||||
try:
|
||||
input_dict = {"samples": a_val}
|
||||
res = node.execute("conv(a, 3, 3, 3, 1.0)", a=input_dict)
|
||||
@@ -240,6 +249,7 @@ def test_conv_3d_asymmetric():
|
||||
print(f"3D Asymmetric Failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
test_conv_1d()
|
||||
@@ -254,5 +264,6 @@ if __name__ == "__main__":
|
||||
print("All Conv tests passed!")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
@@ -8,19 +8,16 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")
|
||||
|
||||
from more_math.helper_functions import eval_tensor_expr, eval_float_expr
|
||||
|
||||
|
||||
def test_all_functions():
|
||||
# Setup some test data
|
||||
a_val = 2.0
|
||||
b_val = 3.0
|
||||
tensor_a = torch.tensor([1.0, 2.0, 3.0])
|
||||
tensor_b = torch.tensor([0.5, 1.5, 2.5])
|
||||
|
||||
variables = {
|
||||
'a': a_val, 'b': b_val,
|
||||
'ta': tensor_a, 'tb': tensor_b,
|
||||
'x': 0.5, 'y': 1.0, 'z': 2.0
|
||||
}
|
||||
|
||||
|
||||
variables = {"a": a_val, "b": b_val, "ta": tensor_a, "tb": tensor_b, "x": 0.5, "y": 1.0, "z": 2.0}
|
||||
|
||||
# helper for assertions
|
||||
def check(expr, expected_scalar=None, vars=variables):
|
||||
# Test scalar
|
||||
@@ -29,7 +26,7 @@ def test_all_functions():
|
||||
if isinstance(res_s, (int, float)):
|
||||
assert abs(res_s - expected_scalar) < 1e-4, f"Scalar {expr} failed: {res_s} != {expected_scalar}"
|
||||
# if expected is tensor we check differently
|
||||
|
||||
|
||||
# Test tensor
|
||||
res_t = eval_tensor_expr(expr, vars, (3,))
|
||||
assert torch.is_tensor(res_t) or isinstance(res_t, (list, int, float))
|
||||
@@ -49,20 +46,20 @@ def test_all_functions():
|
||||
# acosh(1) = 0
|
||||
check("acosh(1)", 0.0)
|
||||
check("atanh(0)", 0.0)
|
||||
|
||||
|
||||
check("abs(-5)", 5.0)
|
||||
check("| -10 |", 10.0) # AbsExp
|
||||
check("| -10 |", 10.0) # AbsExp
|
||||
check("sqrt(16)", 4.0)
|
||||
check("ln(e)", 1.0)
|
||||
check("log(100)", 2.0)
|
||||
check("exp(1)", math.e)
|
||||
|
||||
|
||||
check("floor(1.9)", 1.0)
|
||||
check("ceil(1.1)", 2.0)
|
||||
check("round(1.5)", 2.0)
|
||||
check("gamma(3)", 2.0) # gamma(n) = (n-1)!
|
||||
check("gamma(3)", 2.0) # gamma(n) = (n-1)!
|
||||
check("sigm(0)", 0.5)
|
||||
|
||||
|
||||
check("fract(1.25)", 0.25)
|
||||
check("relu(-5)", 0.0)
|
||||
check("relu(5)", 5.0)
|
||||
@@ -71,16 +68,16 @@ def test_all_functions():
|
||||
check("gelu(0)", 0.0)
|
||||
check("sign(-10)", -1.0)
|
||||
check("sign(10)", 1.0)
|
||||
check("angle(ta)") # test complex angle? no, just ensuring it runs
|
||||
|
||||
check("angle(ta)") # test complex angle? no, just ensuring it runs
|
||||
|
||||
print("--- Testing Two-Arg Functions ---")
|
||||
check("pow(2, 3)", 8.0)
|
||||
check("atan2(1, 1)", math.pi/4)
|
||||
check("atan2(1, 1)", math.pi / 4)
|
||||
check("tmin(5, 10)", 5.0)
|
||||
check("tmax(5, 10)", 10.0)
|
||||
check("step(0.5, 0.2)", 1.0) # step(x, edge) = 1 if x>=edge
|
||||
check("step(0.5, 0.2)", 1.0) # step(x, edge) = 1 if x>=edge
|
||||
check("step(0.1, 0.2)", 0.0)
|
||||
|
||||
|
||||
print("--- Testing Operators ---")
|
||||
check("1 + 2", 3.0)
|
||||
check("5 - 3", 2.0)
|
||||
@@ -88,7 +85,7 @@ def test_all_functions():
|
||||
check("10 / 2", 5.0)
|
||||
check("7 % 3", 1.0)
|
||||
check("2 ^ 3", 8.0)
|
||||
|
||||
|
||||
print("--- Testing Boolean/Comparison ---")
|
||||
check("5 > 3", 1.0)
|
||||
check("5 < 3", 0.0)
|
||||
@@ -96,33 +93,33 @@ def test_all_functions():
|
||||
check("5 <= 4", 0.0)
|
||||
check("2 == 2", 1.0)
|
||||
check("2 != 3", 1.0)
|
||||
|
||||
|
||||
print("--- Testing Ternary/N-ary ---")
|
||||
check("clamp(5, 0, 10)", 5.0)
|
||||
check("clamp(-5, 0, 10)", 0.0)
|
||||
check("lerp(0, 10, 0.5)", 5.0)
|
||||
check("smoothstep(0.5, 0, 1)", 0.5) # smoothstep(x, edge0, edge1)
|
||||
|
||||
check("smoothstep(0.5, 0, 1)", 0.5) # smoothstep(x, edge0, edge1)
|
||||
|
||||
check("smin(1, 2, 3, 0)", 0.0)
|
||||
check("smax(1, 5, 2)", 5.0)
|
||||
|
||||
|
||||
print("--- Testing Tensor Specifics (Norm, Map, Conv, FFT) ---")
|
||||
# Verify abs(x) vs |x|
|
||||
res_abs_func = eval_tensor_expr("abs(ta)", variables, (3,))
|
||||
assert torch.equal(res_abs_func, torch.abs(tensor_a))
|
||||
|
||||
|
||||
res_abs_exp = eval_tensor_expr("|ta|", variables, (3,))
|
||||
assert torch.is_tensor(res_abs_exp) and res_abs_exp.numel() == 1
|
||||
assert torch.allclose(res_abs_exp, torch.linalg.norm(tensor_a))
|
||||
|
||||
check("tnorm(ta)")
|
||||
check("snorm(ta)")
|
||||
check("map(ta, x)") # 1D map
|
||||
check("map(ta, x)") # 1D map
|
||||
# Verify conv with kernel variables kW
|
||||
check("conv(ta, 3, kW)") # 1D conv, size 3, value is kW (which is 3.0)
|
||||
check("conv(ta, 3, kW)") # 1D conv, size 3, value is kW (which is 3.0)
|
||||
check("permute(ta, [0])")
|
||||
check("reshape(ta, [3, 1])")
|
||||
|
||||
|
||||
# FFT/IFFT
|
||||
# We need a shape for FFT usually
|
||||
res_fft = eval_tensor_expr("fft(ta)", variables, (3,))
|
||||
@@ -131,17 +128,17 @@ def test_all_functions():
|
||||
|
||||
# Multi-dim Permute
|
||||
tensor_2d = torch.randn(2, 3)
|
||||
vars_2d = {'t2': tensor_2d}
|
||||
vars_2d = {"t2": tensor_2d}
|
||||
res_perm = eval_tensor_expr("permute(t2, [1, 0])", vars_2d, (3, 2))
|
||||
assert res_perm.shape == (3, 2)
|
||||
|
||||
|
||||
print("--- Testing List and Constants ---")
|
||||
check("[1, 2, 3] + 1")
|
||||
check("pi", math.pi)
|
||||
check("e", math.e)
|
||||
check("print(1)", 1.0)
|
||||
check("pshp(ta)")
|
||||
|
||||
|
||||
print("--- Testing Swap ---")
|
||||
# swap(tensor, dim, i, j)
|
||||
# ta = [1, 2, 3]
|
||||
@@ -151,5 +148,6 @@ def test_all_functions():
|
||||
|
||||
print("All coverage tests passed!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_all_functions()
|
||||
|
||||
+70
-62
@@ -18,6 +18,7 @@ if _comfy_root not in sys.path:
|
||||
import torch
|
||||
from more_math.ModelMathNode import ModelMathNode
|
||||
|
||||
|
||||
class MockModelPatcher:
|
||||
def __init__(self, state_dict):
|
||||
self.model = comfy_api.Model()
|
||||
@@ -35,38 +36,33 @@ class MockModelPatcher:
|
||||
def add_patches(self, patches, strength1, strength2):
|
||||
self.patches.update(patches)
|
||||
|
||||
|
||||
def test_model_math_simple_add():
|
||||
# Setup inputs
|
||||
sd_a = {
|
||||
"layer1.weight": torch.tensor([1.0, 2.0]),
|
||||
"layer2.bias": torch.tensor([0.5])
|
||||
}
|
||||
sd_a = {"layer1.weight": torch.tensor([1.0, 2.0]), "layer2.bias": torch.tensor([0.5])}
|
||||
a = MockModelPatcher(sd_a)
|
||||
|
||||
sd_b = {
|
||||
"layer1.weight": torch.tensor([0.5, 0.5]),
|
||||
"layer2.bias": torch.tensor([0.5])
|
||||
}
|
||||
|
||||
sd_b = {"layer1.weight": torch.tensor([0.5, 0.5]), "layer2.bias": torch.tensor([0.5])}
|
||||
b = MockModelPatcher(sd_b)
|
||||
|
||||
# Expression: a + b
|
||||
expr = "a + b"
|
||||
|
||||
|
||||
# Execute
|
||||
result_tuple = ModelMathNode.execute(expr, a, b=b)
|
||||
result_model = result_tuple[0]
|
||||
|
||||
|
||||
# Verify patches
|
||||
patches = result_model.patches
|
||||
|
||||
# "layer1.weight":
|
||||
|
||||
# "layer1.weight":
|
||||
# a = [1.0, 2.0]
|
||||
# b = [0.5, 0.5]
|
||||
# res = [1.5, 2.5]
|
||||
# diff = res - a = [0.5, 0.5]
|
||||
assert "layer1.weight" in patches
|
||||
assert torch.allclose(patches["layer1.weight"][0], torch.tensor([0.5, 0.5]))
|
||||
|
||||
|
||||
# "layer2.bias":
|
||||
# a = 0.5, b = 0.5
|
||||
# res = 1.0
|
||||
@@ -74,66 +70,69 @@ def test_model_math_simple_add():
|
||||
assert "layer2.bias" in patches
|
||||
assert torch.allclose(patches["layer2.bias"][0], torch.tensor([0.5]))
|
||||
|
||||
|
||||
def test_model_math_zero_default():
|
||||
# Test handling of missing b (should be zero)
|
||||
sd_a = { "w": torch.tensor([10.0]) }
|
||||
sd_a = {"w": torch.tensor([10.0])}
|
||||
a = MockModelPatcher(sd_a)
|
||||
|
||||
|
||||
# expr: a + b. b is None -> b=0.
|
||||
# res = 10 + 0 = 10.
|
||||
# diff = 10 - 10 = 0.
|
||||
# Should skip patch.
|
||||
|
||||
|
||||
result_tuple = ModelMathNode.execute("a + b", a, b=None)
|
||||
result_model = result_tuple[0]
|
||||
|
||||
|
||||
assert len(result_model.patches) == 0
|
||||
|
||||
|
||||
def test_model_math_custom_formula():
|
||||
sd_a = { "w": torch.tensor([2.0]) }
|
||||
sd_a = {"w": torch.tensor([2.0])}
|
||||
a = MockModelPatcher(sd_a)
|
||||
|
||||
|
||||
# expr: a * w
|
||||
# w input = 0.5
|
||||
|
||||
|
||||
result_tuple = ModelMathNode.execute("a * w", a, w=0.5)
|
||||
result_model = result_tuple[0]
|
||||
|
||||
|
||||
# res = 2.0 * 0.5 = 1.0
|
||||
# diff = 1.0 - 2.0 = -1.0
|
||||
patches = result_model.patches
|
||||
assert "w" in patches
|
||||
assert torch.allclose(patches["w"][0], torch.tensor([-1.0]))
|
||||
|
||||
|
||||
def test_model_math_device_mismatch():
|
||||
# Simulating a scenario where model is on a specific device (we can't easily use 'cuda' in CI/CPU env, but we can check if visitor respects it)
|
||||
# We will iterate manually over ModelMathNode logic logic to check TensorEvalVisitor behavior or mock a device.
|
||||
# Actually, simpler: create a tensor on a specific device (if possible) or just ensure visitor uses the device of 'a'.
|
||||
|
||||
|
||||
# Since we might only have CPU, we can't fully reproduce 'cpu vs cuda' crash without cuda.
|
||||
# But we can verify that the scalar created by visitor.visitNumberExp has the same device as 'a'.
|
||||
|
||||
|
||||
from more_math.Parser.TensorEvalVisitor import TensorEvalVisitor
|
||||
from antlr4 import InputStream, CommonTokenStream
|
||||
from more_math.Parser.MathExprLexer import MathExprLexer
|
||||
from more_math.Parser.MathExprParser import MathExprParser
|
||||
|
||||
|
||||
# Mock a tensor with a mocked .device property if we can't use real devices
|
||||
# But torch tensors created with torch.tensor are real.
|
||||
# We will rely on inspection of the visitor code change.
|
||||
# Or strict check:
|
||||
|
||||
# We will rely on inspection of the visitor code change.
|
||||
# Or strict check:
|
||||
|
||||
# Let's try to pass a dummy device string if create_tensor allows, or just check the code path.
|
||||
# Better: Inspect the created tensor from visitor.
|
||||
|
||||
|
||||
tsr = torch.zeros((1,))
|
||||
# We interpret "device mismatch" as: created scalars didn't pick up the device of 'tsr'.
|
||||
# We can force 'tsr' to be on a specific device if available, but likely only 'cpu' is available.
|
||||
# Use a mock object for 'a' that claims to be on 'cuda:0', even if it isn't real tensor?
|
||||
# Use a mock object for 'a' that claims to be on 'cuda:0', even if it isn't real tensor?
|
||||
# No, TensorEvalVisitor expects torch operations.
|
||||
|
||||
|
||||
# We will verify that TensorEvalVisitor creates tensors with the same device as input 'a'.
|
||||
|
||||
|
||||
# Setup
|
||||
expr = "0.5"
|
||||
input_stream = InputStream(expr)
|
||||
@@ -141,108 +140,117 @@ def test_model_math_device_mismatch():
|
||||
stream = CommonTokenStream(lexer)
|
||||
parser = MathExprParser(stream)
|
||||
tree = parser.expr()
|
||||
|
||||
variables = {'a': torch.zeros(1)}
|
||||
|
||||
variables = {"a": torch.zeros(1)}
|
||||
# We want to ensure that if we had a non-cpu device, it would use it.
|
||||
# Since we can't really test this without a GPU, we will write the fix and verify it analytically
|
||||
# Since we can't really test this without a GPU, we will write the fix and verify it analytically
|
||||
# or use a mock that wraps a tensor but intercepts .device?
|
||||
|
||||
|
||||
# Let's just create a test that ensures the result of '0.5' is a tensor that has the same properties as we expect.
|
||||
pass
|
||||
|
||||
|
||||
class MockPatcherContainer:
|
||||
def __init__(self, state_dict):
|
||||
self.patcher = MockModelPatcher(state_dict)
|
||||
|
||||
def clone(self):
|
||||
new_obj = MockPatcherContainer(self.patcher.model.state_dict())
|
||||
new_obj.patcher = self.patcher.clone() # Mock clone behavior
|
||||
new_obj.patcher = self.patcher.clone() # Mock clone behavior
|
||||
return new_obj
|
||||
|
||||
|
||||
def add_patches(self, patches, s1, s2):
|
||||
self.patcher.add_patches(patches, s1, s2)
|
||||
|
||||
|
||||
def test_clip_math():
|
||||
from more_math.ModelMathNode import CLIPMathNode
|
||||
sd_a = { "text_model.encoder.layers.0.mlp.fc1.weight": torch.tensor([1.0]) }
|
||||
|
||||
sd_a = {"text_model.encoder.layers.0.mlp.fc1.weight": torch.tensor([1.0])}
|
||||
clip_a = MockPatcherContainer(sd_a)
|
||||
|
||||
|
||||
result_tuple = CLIPMathNode.execute("a * 2", clip_a)
|
||||
result_clip = result_tuple[0]
|
||||
|
||||
|
||||
# 1.0 * 2 = 2.0. diff = 1.0.
|
||||
patches = result_clip.patcher.patches
|
||||
assert "text_model.encoder.layers.0.mlp.fc1.weight" in patches
|
||||
assert torch.allclose(patches["text_model.encoder.layers.0.mlp.fc1.weight"][0], torch.tensor([1.0]))
|
||||
|
||||
|
||||
def test_vae_math():
|
||||
from more_math.ModelMathNode import VAEMathNode
|
||||
sd_a = { "decoder.conv_in.weight": torch.tensor([1.0]) }
|
||||
|
||||
sd_a = {"decoder.conv_in.weight": torch.tensor([1.0])}
|
||||
# VAE acts like PatcherContainer but no clone() on the VAE itself usually.
|
||||
# But for mock we can use PatcherContainer structure because VAEMathNode expects .patcher
|
||||
|
||||
|
||||
# Mock VAE object
|
||||
class MockVAE:
|
||||
def __init__(self, state_dict):
|
||||
self.patcher = MockModelPatcher(state_dict)
|
||||
|
||||
# No clone method on VAE
|
||||
|
||||
|
||||
vae_a = MockVAE(sd_a)
|
||||
|
||||
|
||||
result_tuple = VAEMathNode.execute("a + 1", vae_a)
|
||||
result_vae = result_tuple[0]
|
||||
|
||||
|
||||
# Check if result is a copy (not same object)
|
||||
assert result_vae is not vae_a
|
||||
# Check if patcher is cloned
|
||||
assert result_vae.patcher is not vae_a.patcher
|
||||
|
||||
|
||||
# 1.0 + 1 = 2.0. diff = 1.0
|
||||
patches = result_vae.patcher.patches
|
||||
assert "decoder.conv_in.weight" in patches
|
||||
assert torch.allclose(patches["decoder.conv_in.weight"][0], torch.tensor([1.0]))
|
||||
|
||||
|
||||
def test_model_math_disjoint_keys():
|
||||
# ModelMathNode uses 'a' as the base.
|
||||
# Keys in 'b' that are NOT in 'a' should be ignored in the loop generally,
|
||||
# ModelMathNode uses 'a' as the base.
|
||||
# Keys in 'b' that are NOT in 'a' should be ignored in the loop generally,
|
||||
# unless logic iterates over union.
|
||||
# Current logic: for key, tens_a in sd_a.items(): ...
|
||||
# So keys only in B are IGNORED.
|
||||
|
||||
sd_a = { "common": torch.tensor([1.0]) }
|
||||
sd_b = { "common": torch.tensor([2.0]), "extra_b": torch.tensor([3.0]) }
|
||||
|
||||
|
||||
sd_a = {"common": torch.tensor([1.0])}
|
||||
sd_b = {"common": torch.tensor([2.0]), "extra_b": torch.tensor([3.0])}
|
||||
|
||||
a = MockPatcherContainer(sd_a)
|
||||
b = MockPatcherContainer(sd_b)
|
||||
|
||||
|
||||
# expr: a + b
|
||||
# common: 1 + 2 = 3. diff = 2.
|
||||
# extra_b: should NOT be in patches because loop is over 'a'
|
||||
|
||||
|
||||
result_tuple = ModelMathNode.execute("a + b", a.patcher, b=b.patcher)
|
||||
patches = result_tuple[0].patches
|
||||
|
||||
|
||||
assert "common" in patches
|
||||
assert torch.allclose(patches["common"][0], torch.tensor([2.0]))
|
||||
assert "extra_b" not in patches
|
||||
|
||||
|
||||
def test_model_math_weighted_merge():
|
||||
# Test w/x/y/z weights
|
||||
sd_a = { "w": torch.tensor([10.0]) }
|
||||
sd_b = { "w": torch.tensor([20.0]) }
|
||||
|
||||
sd_a = {"w": torch.tensor([10.0])}
|
||||
sd_b = {"w": torch.tensor([20.0])}
|
||||
|
||||
a = MockPatcherContainer(sd_a)
|
||||
b = MockPatcherContainer(sd_b)
|
||||
|
||||
|
||||
# expr: lerp(a, b, w) with w=0.5 -> 15.0
|
||||
# diff = 15 - 10 = 5.
|
||||
|
||||
|
||||
result_tuple = ModelMathNode.execute("lerp(a, b, w)", a.patcher, b=b.patcher, w=0.5)
|
||||
patches = result_tuple[0].patches
|
||||
|
||||
|
||||
assert "w" in patches
|
||||
assert torch.allclose(patches["w"][0], torch.tensor([5.0]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_model_math_simple_add()
|
||||
test_model_math_zero_default()
|
||||
|
||||
+35
-20
@@ -32,6 +32,7 @@ from more_math.FloatMathNode import FloatMathNode
|
||||
# Node Initialization and Metadata Tests
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_conditioning_math_node_initialization():
|
||||
node = ConditioningMathNode()
|
||||
assert isinstance(node, ConditioningMathNode)
|
||||
@@ -69,37 +70,33 @@ def test_image_math_node_metadata():
|
||||
# FFT Tests
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_fft_invertibility():
|
||||
# Create random input latent (Batch, Channel, Height, Width)
|
||||
input_tensor = torch.randn(1, 4, 32, 32, dtype=torch.float32)
|
||||
input_dict = {"samples": input_tensor}
|
||||
# Execute ifft(fft(a))
|
||||
result = LatentMathNode.execute(
|
||||
Latent="ifft(fft(a))",
|
||||
a=input_dict
|
||||
)
|
||||
result = LatentMathNode.execute(Latent="ifft(fft(a))", a=input_dict)
|
||||
output_tensor = result[0]["samples"]
|
||||
assert torch.allclose(input_tensor, output_tensor, atol=1e-5), \
|
||||
f"Max difference: {(input_tensor - output_tensor).abs().max()}"
|
||||
assert torch.allclose(input_tensor, output_tensor, atol=1e-5), f"Max difference: {(input_tensor - output_tensor).abs().max()}"
|
||||
|
||||
|
||||
def test_image_fft_dims():
|
||||
# Image input is (Batch, Height, Width, Channel)
|
||||
input_tensor = torch.randn(1, 32, 32, 3, dtype=torch.float32)
|
||||
result = ImageMathNode.execute(
|
||||
Image="ifft(fft(a))",
|
||||
a=input_tensor
|
||||
)
|
||||
result = ImageMathNode.execute(Image="ifft(fft(a))", a=input_tensor)
|
||||
output_tensor = result[0]
|
||||
assert input_tensor.shape == output_tensor.shape
|
||||
assert torch.allclose(input_tensor, output_tensor, atol=1e-5), \
|
||||
assert torch.allclose(input_tensor, output_tensor, atol=1e-5), (
|
||||
f"Image FFT round trip failed. Max diff: {(input_tensor - output_tensor).abs().max()}"
|
||||
)
|
||||
|
||||
|
||||
# ==========================================
|
||||
# Latent Math Basic Functions (Evaluated on Tensors)
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_latent_lerp():
|
||||
node = LatentMathNode()
|
||||
l_a = {"samples": torch.zeros(1, 4, 32, 32)}
|
||||
@@ -111,24 +108,24 @@ def test_latent_lerp():
|
||||
def test_latent_step_true():
|
||||
node = LatentMathNode()
|
||||
# step(0.5, a) where a=0.8 -> 1
|
||||
res_step = node.execute("step(0.5, a)", a={"samples": torch.full((1,1,1,1), 0.8)})[0]["samples"]
|
||||
res_step = node.execute("step(0.5, a)", a={"samples": torch.full((1, 1, 1, 1), 0.8)})[0]["samples"]
|
||||
assert torch.allclose(res_step, torch.ones_like(res_step))
|
||||
|
||||
|
||||
def test_latent_step_false():
|
||||
node = LatentMathNode()
|
||||
# step(0.5, a) where a=0.2 -> 0
|
||||
res_step2 = node.execute("step(0.5, a)", a={"samples": torch.full((1,1,1,1), 0.2)})[0]["samples"]
|
||||
res_step2 = node.execute("step(0.5, a)", a={"samples": torch.full((1, 1, 1, 1), 0.2)})[0]["samples"]
|
||||
assert torch.allclose(res_step2, torch.zeros_like(res_step2))
|
||||
|
||||
|
||||
def test_latent_swap():
|
||||
node = LatentMathNode()
|
||||
t_lat = torch.tensor([0.0, 10.0, 20.0, 30.0]).view(1,4,1,1)
|
||||
t_lat = torch.tensor([0.0, 10.0, 20.0, 30.0]).view(1, 4, 1, 1)
|
||||
# Swap channels 0 and 3 -> 30, 10, 20, 0
|
||||
l_swap = {"samples": t_lat}
|
||||
res_swap = node.execute("swap(a, 1, 0, 3)", a=l_swap)[0]["samples"]
|
||||
expected = torch.tensor([30.0, 10.0, 20.0, 0.0]).view(1,4,1,1)
|
||||
expected = torch.tensor([30.0, 10.0, 20.0, 0.0]).view(1, 4, 1, 1)
|
||||
assert torch.allclose(res_swap, expected)
|
||||
|
||||
|
||||
@@ -157,6 +154,7 @@ def test_latent_fract():
|
||||
# Float Math Basic Functions (Evaluated on Scalars)
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_float_lerp():
|
||||
node = FloatMathNode()
|
||||
res = node.execute("lerp(a, b, 0.5)", a=0.0, b=10.0)[0]
|
||||
@@ -185,6 +183,7 @@ def test_float_smoothstep():
|
||||
# Float Math Extended Functions
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_float_fract():
|
||||
node = FloatMathNode()
|
||||
res = node.execute("fract(a)", a=1.5)[0]
|
||||
@@ -221,6 +220,7 @@ def test_float_gelu():
|
||||
# Latent Math Extended Functions
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_latent_smoothstep():
|
||||
node = LatentMathNode()
|
||||
l_a = {"samples": torch.zeros(1, 4, 32, 32)}
|
||||
@@ -246,6 +246,7 @@ def test_latent_gelu():
|
||||
# Image Math Operations
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_image_lerp():
|
||||
node = ImageMathNode()
|
||||
img_red = torch.tensor([1.0, 0.0, 0.0]).view(1, 1, 1, 3)
|
||||
@@ -267,6 +268,7 @@ def test_image_swap():
|
||||
# Nested Expressions
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_float_nested_expressions_true():
|
||||
node = FloatMathNode()
|
||||
# lerp(0, 10, step(0.5, 0.8)) -> lerp(0, 10, 1) -> 10
|
||||
@@ -285,6 +287,7 @@ def test_float_nested_expressions_false():
|
||||
# 5D Tensor Support
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_5d_tensors_identity():
|
||||
node = LatentMathNode()
|
||||
samples = torch.randn(1, 5, 4, 32, 32)
|
||||
@@ -315,12 +318,14 @@ def test_5d_tensors_fft():
|
||||
# Noise Math Node 5D Support
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_noise_math_5d():
|
||||
from more_math.NoiseMathNode import NoiseMathNode
|
||||
|
||||
class MockNoise:
|
||||
def __init__(self, tensor):
|
||||
self.tensor = tensor
|
||||
|
||||
def generate_noise(self, input_latent):
|
||||
return self.tensor
|
||||
|
||||
@@ -340,6 +345,7 @@ def test_noise_math_5d():
|
||||
# NestedTensor Support
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_nested_tensor_support():
|
||||
try:
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
@@ -354,16 +360,18 @@ def test_nested_tensor_support():
|
||||
|
||||
res_lat = node.execute("a + 1.0", a=l_in)[0]["samples"]
|
||||
|
||||
assert getattr(res_lat, 'is_nested', False)
|
||||
assert getattr(res_lat, "is_nested", False)
|
||||
res_list = res_lat.unbind()
|
||||
assert len(res_list) == 2
|
||||
assert torch.allclose(res_list[0], torch.full_like(t1, 2.0))
|
||||
assert torch.allclose(res_list[1], torch.full_like(t2, 3.0))
|
||||
|
||||
|
||||
# ==========================================
|
||||
# Comprehensive Math Function Tests
|
||||
# ==========================================
|
||||
|
||||
|
||||
def test_trig_functions():
|
||||
node = FloatMathNode()
|
||||
# Sin/Cos checks
|
||||
@@ -373,6 +381,7 @@ def test_trig_functions():
|
||||
# tan(0) = 0
|
||||
assert abs(node.execute("tan(0)", a=0.0)[0] - 0.0) < 1e-5
|
||||
|
||||
|
||||
def test_inverse_trig_functions():
|
||||
node = FloatMathNode()
|
||||
# asin(0) = 0, acos(1) = 0, atan(0) = 0
|
||||
@@ -380,6 +389,7 @@ def test_inverse_trig_functions():
|
||||
assert abs(node.execute("acos(1)", a=0.0)[0] - 0.0) < 1e-5
|
||||
assert abs(node.execute("atan(0)", a=0.0)[0] - 0.0) < 1e-5
|
||||
|
||||
|
||||
def test_pow_log_functions():
|
||||
node = FloatMathNode()
|
||||
# pow(2, 3) = 8
|
||||
@@ -393,15 +403,17 @@ def test_pow_log_functions():
|
||||
# ln(e) = 1. Using 'e' constant logic check or approx 2.718
|
||||
assert abs(node.execute("ln(2.7182818)", a=0.0)[0] - 1.0) < 1e-4
|
||||
|
||||
|
||||
def test_min_max_functions():
|
||||
import sys
|
||||
|
||||
node = FloatMathNode()
|
||||
print("Testing tmin...", flush=True)
|
||||
# tmin(2, 5) = 2, tmax(2, 5) = 5
|
||||
assert abs(node.execute("tmin(2, 5)", a=0.0)[0] - 2.0) < 1e-5
|
||||
print("Testing tmax...", flush=True)
|
||||
assert abs(node.execute("tmax(2, 5)", a=0.0)[0] - 5.0) < 1e-5
|
||||
|
||||
|
||||
# smin/smax (Smooth min/max? Or just multi-arg min/max? TensorEvalVisitor uses stack.min/max)
|
||||
# smin(1, 2, 3) = 1
|
||||
print("Testing smin...", flush=True)
|
||||
@@ -414,6 +426,7 @@ def test_min_max_functions():
|
||||
print(f"smax result: {res_smax} type: {type(res_smax)}", flush=True)
|
||||
assert abs(res_smax - 3.0) < 1e-5
|
||||
|
||||
|
||||
def test_basic_utilities():
|
||||
node = FloatMathNode()
|
||||
# abs(-5) = 5
|
||||
@@ -429,14 +442,17 @@ def test_basic_utilities():
|
||||
assert abs(node.execute("clamp(10, 0, 5)", a=0.0)[0] - 5.0) < 1e-5
|
||||
assert abs(node.execute("clamp(-5, 0, 5)", a=0.0)[0] - 0.0) < 1e-5
|
||||
|
||||
|
||||
def test_advanced_activations():
|
||||
node = FloatMathNode()
|
||||
# sigm(0) = 0.5
|
||||
assert abs(node.execute("sigm(0)", a=0.0)[0] - 0.5) < 1e-5
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
try:
|
||||
|
||||
try:
|
||||
test_conditioning_math_node_initialization()
|
||||
test_conditioning_math_node_metadata()
|
||||
test_latent_math_node_initialization()
|
||||
@@ -452,8 +468,7 @@ if __name__ == "__main__":
|
||||
print("All tests passed!")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
print("All tests in test_more_math.py passed!")
|
||||
|
||||
|
||||
|
||||
+34
-23
@@ -16,146 +16,156 @@ from more_math.Parser.MathExprLexer import MathExprLexer
|
||||
from more_math.Parser.MathExprParser import MathExprParser
|
||||
from antlr4 import InputStream, CommonTokenStream
|
||||
|
||||
|
||||
def parse_and_visit(expr_str, variables):
|
||||
lexer = MathExprLexer(InputStream(expr_str))
|
||||
stream = CommonTokenStream(lexer)
|
||||
parser = MathExprParser(stream)
|
||||
tree = parser.expr()
|
||||
|
||||
|
||||
# We might need to pass shape/device if UnifiedMathVisitor requires it for tensor creation
|
||||
# For now assuming it can infer or defaults.
|
||||
# The original TensorEvalVisitor required shape. Unified might need it for "1.0" -> Tensor promotion cases?
|
||||
# Or maybe "1.0" stays scalar till needed?
|
||||
# Let's assume we pass a default shape/device if needed, but for scalar tests we might not need it.
|
||||
|
||||
shape = (1, 1, 1, 1) # Dummy shape
|
||||
|
||||
shape = (1, 1, 1, 1) # Dummy shape
|
||||
visitor = UnifiedMathVisitor(variables, shape)
|
||||
return visitor.visit(tree)
|
||||
|
||||
|
||||
def test_scalar_ops():
|
||||
vars = {"a": 2.0, "b": 3.0}
|
||||
|
||||
|
||||
assert parse_and_visit("a + b", vars) == 5.0
|
||||
assert parse_and_visit("a * b", vars) == 6.0
|
||||
assert parse_and_visit("sin(0)", vars) == 0.0
|
||||
assert parse_and_visit("smax(a, b)", vars) == 3.0
|
||||
|
||||
|
||||
# Type check - ensure they are python float/int, not tensor
|
||||
res = parse_and_visit("a + b", vars)
|
||||
assert isinstance(res, (float, int))
|
||||
assert isinstance(res, (float, int))
|
||||
|
||||
|
||||
def test_tensor_ops():
|
||||
t1 = torch.tensor([1.0, 2.0])
|
||||
t2 = torch.tensor([3.0, 4.0])
|
||||
vars = {"t1": t1, "t2": t2, "s": 2.0}
|
||||
|
||||
|
||||
# Tensor + Tensor
|
||||
res = parse_and_visit("t1 + t2", vars)
|
||||
assert isinstance(res, torch.Tensor)
|
||||
assert torch.allclose(res, torch.tensor([4.0, 6.0]))
|
||||
|
||||
|
||||
# Tensor + Scalar
|
||||
res2 = parse_and_visit("t1 * s", vars)
|
||||
assert isinstance(res2, torch.Tensor)
|
||||
assert torch.allclose(res2, torch.tensor([2.0, 4.0]))
|
||||
|
||||
|
||||
# Scalar + Tensor
|
||||
res3 = parse_and_visit("s + t2", vars)
|
||||
assert isinstance(res3, torch.Tensor)
|
||||
assert torch.allclose(res3, torch.tensor([5.0, 6.0]))
|
||||
|
||||
|
||||
def test_list_broadcasting():
|
||||
# Feature: List * Tensor -> Stack of Tensors
|
||||
t = torch.ones((2, 2)) # 2x2 ones
|
||||
t = torch.ones((2, 2)) # 2x2 ones
|
||||
l = [1.0, 2.0, 3.0]
|
||||
vars = {"t": t, "l": l}
|
||||
|
||||
|
||||
# l * t should produce a stack of 3 tensors: 1*t, 2*t, 3*t
|
||||
# Expected shape: (3, 2, 2)
|
||||
res = parse_and_visit("l * t", vars)
|
||||
|
||||
|
||||
assert isinstance(res, torch.Tensor)
|
||||
assert res.shape == (3, 2, 2)
|
||||
assert torch.allclose(res[0], t * 1.0)
|
||||
assert torch.allclose(res[1], t * 2.0)
|
||||
assert torch.allclose(res[2], t * 3.0)
|
||||
|
||||
|
||||
def test_list_scalar_mapping():
|
||||
# Feature: List * Scalar -> List of results
|
||||
l = [1.0, 2.0, 3.0]
|
||||
vars = {"l": l}
|
||||
|
||||
|
||||
res = parse_and_visit("l * 2", vars)
|
||||
assert isinstance(res, list)
|
||||
assert res == [2.0, 4.0, 6.0]
|
||||
|
||||
|
||||
def test_func_dispatch():
|
||||
t = torch.tensor([0.0, math.pi/2])
|
||||
t = torch.tensor([0.0, math.pi / 2])
|
||||
vars = {"t": t, "s": 0.0}
|
||||
|
||||
|
||||
# sin(tensor) -> tensor
|
||||
res_t = parse_and_visit("sin(t)", vars)
|
||||
assert isinstance(res_t, torch.Tensor)
|
||||
assert torch.allclose(res_t, torch.tensor([0.0, 1.0]))
|
||||
|
||||
|
||||
# sin(scalar) -> scalar
|
||||
res_s = parse_and_visit("sin(s)", vars)
|
||||
assert isinstance(res_s, float)
|
||||
assert abs(res_s) < 1e-6
|
||||
|
||||
|
||||
# sin(list) -> list
|
||||
l = [0.0, math.pi/2]
|
||||
l = [0.0, math.pi / 2]
|
||||
vars["l"] = l
|
||||
res_l = parse_and_visit("sin(l)", vars)
|
||||
assert isinstance(res_l, list)
|
||||
assert abs(res_l[0]) < 1e-6
|
||||
assert abs(res_l[1] - 1.0) < 1e-6
|
||||
|
||||
|
||||
def test_power_ops():
|
||||
vars = {"a": 2.0, "b": 3.0}
|
||||
# Scalar ^ Scalar
|
||||
assert parse_and_visit("a ^ b", vars) == 8.0
|
||||
|
||||
|
||||
# Tensor ^ Scalar
|
||||
t = torch.tensor([2.0, 3.0])
|
||||
vars["t"] = t
|
||||
res = parse_and_visit("t ^ 2", vars)
|
||||
assert torch.allclose(res, torch.tensor([4.0, 9.0]))
|
||||
|
||||
|
||||
# Scalar ^ Tensor
|
||||
res2 = parse_and_visit("2 ^ t", vars)
|
||||
assert torch.allclose(res2, torch.tensor([4.0, 8.0]))
|
||||
|
||||
|
||||
def test_hyperbolic_trig():
|
||||
vars = {"s": 0.0}
|
||||
assert parse_and_visit("sinh(s)", vars) == 0.0
|
||||
assert parse_and_visit("cosh(s)", vars) == 1.0
|
||||
assert parse_and_visit("tanh(s)", vars) == 0.0
|
||||
|
||||
|
||||
t = torch.tensor([0.0])
|
||||
vars["t"] = t
|
||||
assert torch.allclose(parse_and_visit("sinh(t)", vars), torch.tensor([0.0]))
|
||||
assert torch.allclose(parse_and_visit("cosh(t)", vars), torch.tensor([1.0]))
|
||||
|
||||
|
||||
def test_kernel_coords():
|
||||
# Simulate visitConvFunc context
|
||||
# Usually grid variables are provided by the visitor during visitConvFunc
|
||||
# We can test if they are correctly handled if present in variables
|
||||
grid = torch.linspace(-1, 1, 3)
|
||||
vars = {"kx": grid, "ky": grid}
|
||||
|
||||
|
||||
# Test if expression using coordinates works
|
||||
res = parse_and_visit("kx^2 + ky^2", vars)
|
||||
assert isinstance(res, torch.Tensor)
|
||||
assert res.shape == grid.shape
|
||||
assert torch.allclose(res, grid**2 + grid**2)
|
||||
|
||||
|
||||
def test_bool_ops():
|
||||
vars = {"a": 1, "b": 0}
|
||||
# Scalar bool
|
||||
assert parse_and_visit("a > b", vars) == 1
|
||||
assert parse_and_visit("a < b", vars) == 0
|
||||
|
||||
|
||||
# Tensor bool
|
||||
t1 = torch.tensor([1.0, 0.0])
|
||||
t2 = torch.tensor([0.0, 1.0])
|
||||
@@ -179,5 +189,6 @@ if __name__ == "__main__":
|
||||
print("All UnifiedMathVisitor tests passed!")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
Reference in New Issue
Block a user