93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
#!/usr/bin/env python
|
|
|
|
"""Tests for `more_math` package."""
|
|
|
|
import unittest
|
|
import torch
|
|
from more_math.ConditioningMathNode import ConditioningMathNode
|
|
from more_math.LatentMathNode import LatentMathNode
|
|
from more_math.ImageMathNode import ImageMathNode
|
|
import tokenize
|
|
from io import StringIO
|
|
|
|
def tokenize_expression(expr):
|
|
if not expr.endswith('\n'):
|
|
expr = expr + '\n'
|
|
f = StringIO(expr)
|
|
tokens = tokenize.generate_tokens(f.readline)
|
|
filtered_tokens = []
|
|
for toktype, tokval, _, _, _ in tokens:
|
|
token_name = tokenize.tok_name[toktype]
|
|
# Opravíme ERRORTOKEN na OP pro !, &, |, ^
|
|
if token_name == 'ERRORTOKEN' and tokval in {'!', '&', '|', '^'}:
|
|
token_name = 'OP'
|
|
if token_name in {'COMMENT', 'NL', 'NEWLINE', 'INDENT', 'DEDENT'}:
|
|
continue
|
|
filtered_tokens.append((token_name, tokval.strip()))
|
|
return filtered_tokens
|
|
|
|
|
|
class TestMoreMath(unittest.TestCase):
|
|
def test_conditioning_math_node_initialization(self):
|
|
node = ConditioningMathNode()
|
|
self.assertIsInstance(node, ConditioningMathNode)
|
|
|
|
def test_conditioning_math_node_metadata(self):
|
|
self.assertEqual(ConditioningMathNode.RETURN_TYPES, ["CONDITIONING"])
|
|
self.assertEqual(ConditioningMathNode.FUNCTION, "EXECUTE_NORMALIZED")
|
|
self.assertEqual(ConditioningMathNode.CATEGORY, "More math")
|
|
|
|
def test_latent_math_node_initialization(self):
|
|
node = LatentMathNode()
|
|
self.assertIsInstance(node, LatentMathNode)
|
|
|
|
def test_latent_math_node_metadata(self):
|
|
self.assertEqual(LatentMathNode.RETURN_TYPES, ["LATENT"])
|
|
self.assertEqual(LatentMathNode.FUNCTION, "EXECUTE_NORMALIZED")
|
|
self.assertEqual(LatentMathNode.CATEGORY, "More math")
|
|
|
|
def test_image_math_node_initialization(self):
|
|
node = ImageMathNode()
|
|
self.assertIsInstance(node, ImageMathNode)
|
|
|
|
def test_image_math_node_metadata(self):
|
|
self.assertEqual(ImageMathNode.RETURN_TYPES, ["IMAGE"])
|
|
self.assertEqual(ImageMathNode.FUNCTION, "EXECUTE_NORMALIZED")
|
|
self.assertEqual(ImageMathNode.CATEGORY, "More math")
|
|
|
|
def test_fft_invertibility(self):
|
|
# 1. Create random input latent (Batch, Channel, Height, Width)
|
|
input_tensor = torch.randn(1, 4, 32, 32, dtype=torch.float32)
|
|
input_dict = {"samples": input_tensor}
|
|
|
|
# 2. Execute ifft(fft(a))
|
|
# Note: execute is a classmethod
|
|
result = LatentMathNode.execute(
|
|
Latent="ifft(fft(a))",
|
|
a=input_dict
|
|
)
|
|
|
|
# 3. Get output tensor
|
|
output_tensor = result[0]["samples"]
|
|
|
|
# 4. Check correctness (approximate equality)
|
|
self.assertTrue(torch.allclose(input_tensor, output_tensor, atol=1e-5), \
|
|
f"Max difference: {(input_tensor - output_tensor).abs().max()}")
|
|
|
|
def test_image_fft_dims(self):
|
|
# Image input is (Batch, Height, Width, Channel)
|
|
# Verify 2D FFT works by doing a round trip
|
|
input_tensor = torch.randn(1, 32, 32, 3, dtype=torch.float32)
|
|
|
|
result = ImageMathNode.execute(
|
|
Image="ifft(fft(a))",
|
|
a=input_tensor
|
|
)
|
|
output_tensor = result[0]
|
|
|
|
self.assertEqual(input_tensor.shape, output_tensor.shape)
|
|
self.assertTrue(torch.allclose(input_tensor, output_tensor, atol=1e-5), \
|
|
f"Image FFT round trip failed. Max diff: {(input_tensor - output_tensor).abs().max()}")
|
|
|
|
|