47 lines
2.2 KiB
Python
47 lines
2.2 KiB
Python
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)
|