Merge 11c1d6cc28 into 1d1099bea0

comfy_extras/nodes_rebatch.py

@@ -27,8 +27,8 @@ class LatentRebatch(io.ComfyNode):
         samples = latents[list_ind]['samples']
         shape = samples.shape
         mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')
-        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:
-            torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
+        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2] * 8:
+            mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
         if mask.shape[0] < samples.shape[0]:
             mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
         if 'batch_index' in latents[list_ind]:

tests-unit/comfy_extras_test/nodes_rebatch_test.py

@@ -0,0 +1,46 @@
+from unittest.mock import patch, MagicMock
+
+import torch
+
+mock_nodes = MagicMock()
+mock_nodes.MAX_RESOLUTION = 16384
+mock_server = MagicMock()
+
+with patch.dict("sys.modules", {"nodes": mock_nodes, "server": mock_server}):
+    from comfy_extras.nodes_rebatch import LatentRebatch
+
+
+class TestLatentRebatchGetBatch:
+    def test_default_mask_matches_pixel_resolution(self):
+        # a latent without a noise_mask gets an all-ones mask at samples * 8
+        latents = [{"samples": torch.zeros(1, 4, 16, 16)}]
+        _, mask, _ = LatentRebatch.get_batch(latents, 0, 0)
+        assert mask.shape == (1, 1, 128, 128)
+
+    def test_matching_mask_is_kept(self):
+        latents = [{"samples": torch.zeros(1, 4, 16, 16), "noise_mask": torch.ones(1, 1, 128, 128)}]
+        _, mask, _ = LatentRebatch.get_batch(latents, 0, 0)
+        assert mask.shape == (1, 1, 128, 128)
+
+    def test_mismatched_mask_is_resized_to_pixel_resolution(self):
+        # SetLatentNoiseMask stores masks without resizing, so a mask can arrive
+        # at a resolution that does not match samples * 8 and must be scaled up.
+        latents = [{"samples": torch.zeros(1, 4, 16, 16), "noise_mask": torch.ones(1, 1, 48, 48)}]
+        _, mask, _ = LatentRebatch.get_batch(latents, 0, 0)
+        assert mask.shape == (1, 1, 128, 128)
+
+    def test_mismatched_height_only_is_resized(self):
+        latents = [{"samples": torch.zeros(1, 4, 16, 16), "noise_mask": torch.ones(1, 1, 48, 128)}]
+        _, mask, _ = LatentRebatch.get_batch(latents, 0, 0)
+        assert mask.shape == (1, 1, 128, 128)
+
+    def test_batches_with_mismatched_mask_can_be_concatenated(self):
+        # a resized mask must line up with another latent's default mask so the
+        # two batches can be concatenated instead of raising a size mismatch.
+        with_mask = [{"samples": torch.zeros(1, 4, 16, 16), "noise_mask": torch.ones(1, 1, 48, 48)}]
+        without_mask = [{"samples": torch.zeros(1, 4, 16, 16)}]
+        batch_a = LatentRebatch.get_batch(with_mask, 0, 0)
+        batch_b = LatentRebatch.get_batch(without_mask, 0, 1)
+        samples, mask, _ = LatentRebatch.cat_batch(batch_a, batch_b)
+        assert samples.shape == (2, 4, 16, 16)
+        assert mask.shape == (2, 1, 128, 128)