From 11c1d6cc28c780bb80357bb95e4adf98206b574d Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Sun, 12 Jul 2026 01:07:32 +0530 Subject: [PATCH] Fix RebatchLatents dropping resized noise mask get_batch() called torch.nn.functional.interpolate on a mismatched noise mask but discarded the result, so the mask kept its original size. The height check was also missing the *8 pixel scale. When a rebatched latent carried a noise mask that did not match samples * 8 (e.g. from SetLatentNoiseMask, which stores masks unscaled), the unresized mask later failed to concatenate with another latent's default mask and raised a size-mismatch RuntimeError. Assign the interpolate result back to mask and compare the height against shape[-2] * 8. --- comfy_extras/nodes_rebatch.py | 4 +- .../comfy_extras_test/nodes_rebatch_test.py | 46 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) create mode 100644 tests-unit/comfy_extras_test/nodes_rebatch_test.py diff --git a/comfy_extras/nodes_rebatch.py b/comfy_extras/nodes_rebatch.py index 2185385f0..3d55101a9 100644 --- a/comfy_extras/nodes_rebatch.py +++ b/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]: diff --git a/tests-unit/comfy_extras_test/nodes_rebatch_test.py b/tests-unit/comfy_extras_test/nodes_rebatch_test.py new file mode 100644 index 000000000..a7e5230e1 --- /dev/null +++ b/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)