102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
import torch
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
from more_math.GuiderMathNode import MathGuider
|
|
|
|
class MockGuider:
|
|
def __init__(self, value, device="cpu"):
|
|
self.value = value
|
|
self.device = device
|
|
self.model_patcher = MagicMock()
|
|
self.model_patcher.load_device = device
|
|
self.model_patcher.offload_device = "cpu"
|
|
self.model_patcher.model_dtype = lambda: torch.float32
|
|
self.original_conds = {}
|
|
self.model_options = {}
|
|
|
|
def __call__(self, x, sigma, model_options={}, seed=None):
|
|
return torch.full_like(x, self.value)
|
|
|
|
class TestMathGuider(unittest.TestCase):
|
|
def test_math_guider_call(self):
|
|
# Setup input guiders
|
|
g0 = MockGuider(1.0)
|
|
g1 = MockGuider(2.0)
|
|
V = {"V0": g0, "V1": g1}
|
|
F = {"F0": 0.5}
|
|
|
|
# Expression: Average V0 and V1
|
|
expr = "V0 * 0.5 + V1 * 0.5"
|
|
|
|
math_guider = MathGuider(V, F, expr)
|
|
|
|
# Pseudo input
|
|
x = torch.zeros((1, 4, 16, 16))
|
|
sigma = torch.tensor(1.0)
|
|
|
|
# Call
|
|
result = math_guider(x, sigma)
|
|
|
|
# Expected: 1.0 * 0.5 + 2.0 * 0.5 = 1.5
|
|
self.assertTrue(torch.allclose(result, torch.tensor(1.5)))
|
|
|
|
def test_math_guider_aliases(self):
|
|
g0 = MockGuider(10.0)
|
|
V = {"V0": g0}
|
|
F = {"F0": 2.0}
|
|
|
|
# a = V0, w = F0
|
|
expr = "a + w"
|
|
|
|
math_guider = MathGuider(V, F, expr)
|
|
x = torch.zeros((1, 4, 8, 8))
|
|
sigma = torch.tensor(1.0)
|
|
|
|
result = math_guider(x, sigma)
|
|
self.assertTrue(torch.allclose(result, torch.tensor(12.0)))
|
|
|
|
def test_math_guider_model_patcher(self):
|
|
# Verify that math_guider exposes model_patcher from its input guider
|
|
g0 = MockGuider(1.0)
|
|
V = {"V0": g0}
|
|
F = {}
|
|
math_guider = MathGuider(V, F, "V0")
|
|
|
|
# Check if the property exists and matches g0's patcher
|
|
self.assertIsNotNone(math_guider.model_patcher)
|
|
self.assertEqual(math_guider.model_patcher, g0.model_patcher)
|
|
|
|
def test_math_guider_model_patcher_missing(self):
|
|
# Verify behavior when input guiders don't have model_patcher (e.g. None or broken)
|
|
g0 = MockGuider(1.0)
|
|
del g0.model_patcher # force remove
|
|
V = {"V0": g0}
|
|
F = {}
|
|
math_guider = MathGuider(V, F, "V0")
|
|
self.assertIsNone(math_guider.model_patcher)
|
|
|
|
def test_math_guider_steps_context(self):
|
|
# Mock sigmas: [10.0, 5.0, 0.0] -> 2 steps
|
|
sigmas = torch.tensor([10.0, 5.0, 0.0])
|
|
g0 = MockGuider(1.0)
|
|
V = {"V0": g0}
|
|
|
|
math_guider = MathGuider(V, {}, "current_step / steps")
|
|
math_guider.sigmas = sigmas # sets sigmas directly for testing
|
|
|
|
# Step 0: sigma = 10.0
|
|
x = torch.zeros((1, 1, 1, 1))
|
|
res0 = math_guider(x, torch.tensor(10.0))
|
|
self.assertTrue(torch.allclose(res0, torch.tensor(0.0 / 2.0)))
|
|
|
|
# Step 1: sigma = 5.0
|
|
res1 = math_guider(x, torch.tensor(5.0))
|
|
self.assertTrue(torch.allclose(res1, torch.tensor(1.0 / 2.0)))
|
|
|
|
# Intermediate sigma should find closest
|
|
res_near = math_guider(x, torch.tensor(4.8))
|
|
self.assertTrue(torch.allclose(res_near, torch.tensor(1.0 / 2.0)))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|