54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from upscale.conditioning import clone_conditioning, clone_control_chain
|
|
|
|
|
|
class _FakeControl:
|
|
def __init__(self, hint, previous=None):
|
|
self.cond_hint_original = hint
|
|
self.previous_controlnet = previous
|
|
|
|
|
|
class ConditioningUtilsTests(unittest.TestCase):
|
|
def test_clone_control_chain_clones_hints_and_links(self):
|
|
tail = _FakeControl(torch.ones((1, 2)))
|
|
head = _FakeControl(torch.zeros((1, 2)), previous=tail)
|
|
|
|
cloned = clone_control_chain(head, clone_hint=True)
|
|
|
|
self.assertIsNot(cloned, head)
|
|
self.assertIsNot(cloned.previous_controlnet, tail)
|
|
self.assertTrue(torch.equal(cloned.cond_hint_original, head.cond_hint_original))
|
|
self.assertTrue(torch.equal(cloned.previous_controlnet.cond_hint_original, tail.cond_hint_original))
|
|
self.assertIsNot(cloned.cond_hint_original, head.cond_hint_original)
|
|
|
|
def test_clone_conditioning_clones_emb_and_known_fields(self):
|
|
control = _FakeControl(torch.ones((1, 3)))
|
|
cond = [
|
|
[
|
|
torch.tensor([[1.0, 2.0, 3.0]]),
|
|
{
|
|
"control": control,
|
|
"mask": torch.ones((1, 1, 1)),
|
|
"pooled_output": torch.zeros((1, 4)),
|
|
"area": [0, 1, 2, 3],
|
|
},
|
|
]
|
|
]
|
|
|
|
cloned = clone_conditioning(cond, clone_hints=True)
|
|
|
|
self.assertTrue(torch.equal(cloned[0][0], cond[0][0]))
|
|
self.assertIsNot(cloned[0][0], cond[0][0])
|
|
self.assertIsNot(cloned[0][1]["control"], cond[0][1]["control"])
|
|
self.assertIsNot(cloned[0][1]["mask"], cond[0][1]["mask"])
|
|
self.assertIsNot(cloned[0][1]["pooled_output"], cond[0][1]["pooled_output"])
|
|
self.assertEqual(cloned[0][1]["area"], cond[0][1]["area"])
|
|
self.assertIsNot(cloned[0][1]["area"], cond[0][1]["area"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|