35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
|
|
import sys
|
|
import os
|
|
import torch
|
|
sys.path.insert(0, os.path.abspath('src'))
|
|
# Add ComfyUI path for dependencies if needed
|
|
sys.path.insert(0, os.path.abspath('../../'))
|
|
from more_math.LatentMathNode import LatentMathNode
|
|
|
|
print("Testing 5D tensors...")
|
|
node = LatentMathNode()
|
|
# Shape: (1, 5, 4, 32, 32)
|
|
samples = torch.randn(1, 5, 4, 32, 32)
|
|
l_in = {"samples": samples}
|
|
|
|
try:
|
|
print("Test 1: Identity")
|
|
res = node.execute("a * 1.0", a=l_in)[0]["samples"]
|
|
print("Shape:", res.shape)
|
|
assert res.shape == (1, 5, 4, 32, 32)
|
|
|
|
res_t = node.execute("a + T", a=l_in)[0]["samples"]
|
|
print("Shape T:", res_t.shape)
|
|
assert torch.allclose(res_t, samples + 5.0), "T addition failed value check"
|
|
print("T addition passed strict check")
|
|
|
|
print("Test 3: FFT")
|
|
res_fft = node.execute("ifft(fft(a))", a=l_in)[0]["samples"]
|
|
assert torch.allclose(res_fft, samples, atol=1e-5), "FFT roundtrip failed"
|
|
print("FFT passed strict check")
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|