Simplify no_control

2024-06-28 15:52:53 +02:00
parent ddbeb36d52
commit 183d0b2707
2 changed files with 23 additions and 23 deletions
@@ -0,0 +1,19 @@
+from unittest import mock
+
+from coreml_suite.controlnet import no_control
+
+
+def test_no_control():
+    model = mock.Mock()
+    model.expected_inputs = {
+        "additional_residual_0": {"shape": (2, 2, 2)},
+        "additional_residual_1": {"shape": (2, 4, 4)},
+        "additional_residual_2": {"shape": (2, 8, 8)},
+    }
+
+    residual_kwargs = no_control(model)
+
+    assert len(residual_kwargs) == 3
+    assert residual_kwargs["additional_residual_0"].shape == (2, 2, 2)
+    assert residual_kwargs["additional_residual_1"].shape == (2, 4, 4)
+    assert residual_kwargs["additional_residual_2"].shape == (2, 8, 8)