diff --git a/comfy_extras/nodes_latent.py b/comfy_extras/nodes_latent.py index 1f93e34d6..91b5e7bdb 100644 --- a/comfy_extras/nodes_latent.py +++ b/comfy_extras/nodes_latent.py @@ -120,14 +120,14 @@ class LatentInterpolate(io.ComfyNode): s2 = reshape_latent_to(s1.shape, s2) - m1 = torch.linalg.vector_norm(s1, dim=(1)) - m2 = torch.linalg.vector_norm(s2, dim=(1)) + m1 = torch.linalg.vector_norm(s1, dim=1, keepdim=True) + m2 = torch.linalg.vector_norm(s2, dim=1, keepdim=True) s1 = torch.nan_to_num(s1 / m1) s2 = torch.nan_to_num(s2 / m2) t = (s1 * ratio + s2 * (1.0 - ratio)) - mt = torch.linalg.vector_norm(t, dim=(1)) + mt = torch.linalg.vector_norm(t, dim=1, keepdim=True) st = torch.nan_to_num(t / mt) samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio)) diff --git a/tests-unit/comfy_extras_test/nodes_latent_test.py b/tests-unit/comfy_extras_test/nodes_latent_test.py new file mode 100644 index 000000000..001bfe55e --- /dev/null +++ b/tests-unit/comfy_extras_test/nodes_latent_test.py @@ -0,0 +1,42 @@ +from unittest.mock import patch, MagicMock + +import torch + +mock_nodes = MagicMock() +mock_nodes.MAX_RESOLUTION = 16384 +mock_server = MagicMock() + +# comfy.model_management initializes the torch device at import and requires CUDA +# (or --cpu), which the unit-test environment does not provide; stub it for the import. +patched_modules = {"nodes": mock_nodes, "server": mock_server, "comfy.model_management": MagicMock()} +with patch.dict("sys.modules", patched_modules): + # imported first to satisfy the nodes_post_processing <-> nodes_latent circular import + import comfy_extras.nodes_post_processing # noqa: F401 + from comfy_extras.nodes_latent import LatentInterpolate + + +class TestLatentInterpolate: + def test_batch_size_one(self): + s = {"samples": torch.randn(1, 4, 16, 16)} + out = LatentInterpolate.execute(s, s, 0.5)[0] + assert out["samples"].shape == (1, 4, 16, 16) + + def test_batched_latent_does_not_crash(self): + s1 = {"samples": torch.randn(2, 4, 16, 16)} + s2 = {"samples": torch.randn(2, 4, 16, 16)} + out = LatentInterpolate.execute(s1, s2, 0.5)[0] + assert out["samples"].shape == (2, 4, 16, 16) + + def test_each_batch_element_is_interpolated_independently(self): + # normalization must be per latent, not mixed across the batch. batch == channels + # (4) is the case that does not crash but silently corrupts without the fix. + a = torch.randn(1, 4, 8, 8) + b = torch.randn(1, 4, 8, 8) + filler = torch.randn(3, 4, 8, 8) + single = LatentInterpolate.execute({"samples": a}, {"samples": b}, 0.3)[0]["samples"] + batched = LatentInterpolate.execute( + {"samples": torch.cat([a, filler])}, + {"samples": torch.cat([b, filler])}, + 0.3, + )[0]["samples"] + assert torch.allclose(single[0], batched[0], atol=1e-5)