img2img works

tests/test_chunks.py

@@ -7,7 +7,11 @@ import torch
 from comfy.model_management import get_torch_device
 from coreml_suite.latents import chunk_batch, merge_chunks
 from coreml_suite.controlnet import chunk_control
-from coreml_suite.models import CoreMLModelWrapper, get_model_config
+from coreml_suite.models import (
+    CoreMLModelWrapper,
+    get_model_config,
+    CoreMLModelWrapperLCM,
+)
 
 
 @pytest.fixture
@@ -16,6 +20,7 @@ def coreml_model():
     model.expected_inputs = {
         "sample": {"shape": (2, 4, 64, 64)},
         "timestep": {"shape": (2,)},
+        "timestep_cond": {"shape": (2, 256)},
         "encoder_hidden_states": {"shape": (2, 768, 1, 77)},
         "additional_residual_0": {"shape": (2, 320, 64, 64)},
         "additional_residual_1": {"shape": (2, 640, 32, 32)},
@@ -54,6 +59,22 @@ def test_merge_chunks(batch_size):
     assert torch.equal(input_tensor, merged)
 
 
+@pytest.fixture
+def inputs():
+    x = torch.randn(1, 4, 64, 64).to(get_torch_device())
+    t = torch.randn([1]).to(get_torch_device())
+    c_crossattn = torch.randn(1, 77, 768).to(get_torch_device())
+    control = {
+        "output": [
+            torch.randn(1, 320, 64, 64).to(get_torch_device()),
+            torch.randn(1, 640, 32, 32).to(get_torch_device()),
+        ],
+    }
+    timestep_cond = torch.randn(1, 256).to(get_torch_device())
+
+    return x, t, c_crossattn, control, timestep_cond
+
+
 @pytest.mark.parametrize(
     "b, target_size, num_chunks",
     [
@@ -95,29 +116,22 @@ def test_chunking_no_control():
     assert chunked == [None, None]
 
 
-def test_chunking_inputs(coreml_model, model_config):
+def test_chunking_inputs(coreml_model, model_config, inputs):
     model = CoreMLModelWrapper(model_config, coreml_model)
-    x = torch.randn(1, 4, 64, 64).to(get_torch_device())
-    t = torch.randn([1]).to(get_torch_device())
-    c_crossattn = torch.randn(1, 77, 768).to(get_torch_device())
-    control = {
-        "output": [
-            torch.randn(1, 320, 64, 64).to(get_torch_device()),
-            torch.randn(1, 640, 32, 32).to(get_torch_device()),
-        ],
-    }
 
-    chunked_x, ts, chunked_context, chunked_control = model.chunk_inputs(
-        x, t, c_crossattn, control
+    chunked_x, ts, chunked_context, chunked_cn, chunked_ts_cond = model.chunk_inputs(
+        *inputs
     )
 
     assert len(chunked_x) == 1
     assert len(ts) == 1
     assert len(chunked_context) == 1
-    assert len(chunked_control) == 1
+    assert len(chunked_cn) == 1
+    assert len(chunked_ts_cond) == 1
 
     assert chunked_x[0].shape == (2, 4, 64, 64)
     assert ts[0].shape == (2,)
     assert chunked_context[0].shape == (2, 77, 768)
-    assert chunked_control[0]["output"][0].shape == (2, 320, 64, 64)
-    assert chunked_control[0]["output"][1].shape == (2, 640, 32, 32)
+    assert chunked_cn[0]["output"][0].shape == (2, 320, 64, 64)
+    assert chunked_cn[0]["output"][1].shape == (2, 640, 32, 32)
+    assert chunked_ts_cond[0].shape == (2, 256)