109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
import torch
|
|
|
|
from mask_ops import (
|
|
mask_row_targets,
|
|
mixed_mod_row,
|
|
quantize_h3_pixel_mask,
|
|
resize_mask_to_canvas,
|
|
snap_audio_mask,
|
|
snap_video_mask,
|
|
)
|
|
|
|
|
|
def test_video_mask_snaps_each_2x2_patch_to_generate():
|
|
mask = torch.zeros((1, 4, 1, 4, 4))
|
|
mask[:, :, :, 1, 1] = 1.0
|
|
snapped = snap_video_mask(mask)
|
|
assert torch.all(snapped[..., :2, :2] == 1)
|
|
assert torch.all(snapped[..., 2:, :] == 0)
|
|
assert torch.all(snapped[..., :, 2:] == 0)
|
|
|
|
|
|
def test_video_mask_handles_odd_latent_axes():
|
|
mask = torch.zeros((1, 2, 1, 3, 5))
|
|
mask[..., 2, 4] = 1.0
|
|
snapped = snap_video_mask(mask)
|
|
assert snapped.shape == mask.shape
|
|
assert torch.all(snapped[..., 2, 4] == 1)
|
|
|
|
|
|
def test_audio_mask_is_binary_and_shared_across_channels():
|
|
mask = torch.zeros((1, 3, 2, 4))
|
|
mask[:, 2, 1, 3] = 0.75
|
|
snapped = snap_audio_mask(mask)
|
|
assert torch.all(snapped[:, :, 1, 3] == 1)
|
|
assert torch.all(snapped[:, :, 0, :] == 0)
|
|
|
|
|
|
def test_mask_row_targets_returns_none_when_everything_generates():
|
|
mask = torch.ones((2, 4, 4))
|
|
assert mask_row_targets(mask, 2, 4, 4) is None
|
|
|
|
|
|
def test_mask_row_targets_matches_patch_order():
|
|
mask = torch.zeros((1, 4, 4))
|
|
mask[:, :2, 2:] = 1.0
|
|
targets = mask_row_targets(mask, 1, 4, 4)
|
|
assert targets.tolist() == [False, True, False, False]
|
|
|
|
|
|
def test_mixed_mod_row_blends_per_token():
|
|
values = torch.tensor([[0.0, 0.0], [10.0, 20.0]])
|
|
weights = torch.tensor([[0.0], [1.0], [0.25]])
|
|
result = mixed_mod_row(values, (1, 0, weights), torch.float32)
|
|
assert torch.allclose(result, torch.tensor([[0.0, 0.0], [10.0, 20.0], [2.5, 5.0]]))
|
|
|
|
|
|
def test_resize_mask_broadcasts_one_frame_across_video():
|
|
mask = torch.zeros((1, 8, 8))
|
|
mask[:, 2:6, 2:6] = 1.0
|
|
resized = resize_mask_to_canvas(mask, frames=3, height=8, width=8)
|
|
assert resized.shape == (3, 1, 8, 8)
|
|
assert torch.allclose(resized[0], resized[2])
|
|
|
|
|
|
def test_runtime_grid_selects_complete_32_pixel_cell():
|
|
mask = torch.zeros((1, 64, 64))
|
|
mask[:, :16, :16] = 1.0
|
|
_, snapped, cells = quantize_h3_pixel_mask(mask, 1, 64, 64)
|
|
assert cells.shape == (1, 1, 2, 2)
|
|
assert cells[0, 0].tolist() == [[1.0, 0.0], [0.0, 0.0]]
|
|
assert torch.all(snapped[:, :32, :32] == 1)
|
|
assert torch.all(snapped[:, 32:, :] == 0)
|
|
|
|
|
|
def test_pixel_coverage_modes_are_conservative():
|
|
mask = torch.zeros((1, 64, 64))
|
|
mask[:, :8, :8] = 1.0
|
|
_, any_mask, _ = quantize_h3_pixel_mask(
|
|
mask, 1, 64, 64, mode="any pixel coverage"
|
|
)
|
|
_, half_mask, _ = quantize_h3_pixel_mask(
|
|
mask, 1, 64, 64, mode="50% pixel coverage"
|
|
)
|
|
assert any_mask.sum() == 32 * 32
|
|
assert half_mask.sum() == 0
|
|
|
|
|
|
def test_cell_growth_and_shrink_are_whole_cells():
|
|
mask = torch.zeros((1, 96, 96))
|
|
mask[:, 32:64, 32:64] = 1.0
|
|
_, grown, _ = quantize_h3_pixel_mask(
|
|
mask, 1, 96, 96, mode="any pixel coverage", cell_adjust=1
|
|
)
|
|
_, shrunk, _ = quantize_h3_pixel_mask(
|
|
mask, 1, 96, 96, mode="any pixel coverage", cell_adjust=-1
|
|
)
|
|
assert grown.sum() == 96 * 96
|
|
assert shrunk.sum() == 0
|
|
|
|
|
|
def test_grid_preview_rejects_non_h3_canvas():
|
|
mask = torch.ones((1, 64, 64))
|
|
try:
|
|
quantize_h3_pixel_mask(mask, 1, 63, 64)
|
|
except ValueError as error:
|
|
assert "divisible by 32" in str(error)
|
|
else:
|
|
raise AssertionError("Expected non-H3 canvas dimensions to be rejected")
|