Fix LatentInterpolate corrupting batched latents
vector_norm(dim=1) dropped the channel axis, so dividing the (B,C,H,W) latent by the (B,H,W) norm broadcast the batch axis onto the channels. With batch_size != channels this raised a RuntimeError, and when they happened to be equal it silently mixed the per-sample normalization. Keep the channel dimension with keepdim=True so each latent is normalized independently.
This commit is contained in:
parent
f3a36e7484
commit
68f5b543d9
|
|
@ -120,14 +120,14 @@ class LatentInterpolate(io.ComfyNode):
|
||||||
|
|
||||||
s2 = reshape_latent_to(s1.shape, s2)
|
s2 = reshape_latent_to(s1.shape, s2)
|
||||||
|
|
||||||
m1 = torch.linalg.vector_norm(s1, dim=(1))
|
m1 = torch.linalg.vector_norm(s1, dim=1, keepdim=True)
|
||||||
m2 = torch.linalg.vector_norm(s2, dim=(1))
|
m2 = torch.linalg.vector_norm(s2, dim=1, keepdim=True)
|
||||||
|
|
||||||
s1 = torch.nan_to_num(s1 / m1)
|
s1 = torch.nan_to_num(s1 / m1)
|
||||||
s2 = torch.nan_to_num(s2 / m2)
|
s2 = torch.nan_to_num(s2 / m2)
|
||||||
|
|
||||||
t = (s1 * ratio + s2 * (1.0 - ratio))
|
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)
|
st = torch.nan_to_num(t / mt)
|
||||||
|
|
||||||
samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio))
|
samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio))
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
Loading…
Reference in New Issue