| import torch |
| from dememwm_import_helper import install_dememwm_namespace |
|
|
| install_dememwm_namespace() |
| from algorithms.worldmem.dememwm.cache import StreamingCache |
| from algorithms.worldmem.dememwm.compression import CausalConv3DDynamicCompressor |
|
|
|
|
| def small_compressor(**kwargs): |
| return CausalConv3DDynamicCompressor( |
| latent_channels=3, |
| dit_hidden_size=8, |
| patch_size=2, |
| conv_kernel_t=3, |
| conv_stride_t=2, |
| max_source_frames=4, |
| **kwargs, |
| ) |
|
|
|
|
| def test_dynamic_compressor_shapes_and_budget(): |
| comp = small_compressor(exclude_latest_local_frames=0) |
| latents = torch.randn(4, 2, 3, 2, 2) |
| frame_indices = torch.arange(4)[:, None].repeat(1, 2) |
| target = torch.tensor([[1, 2], [4, 4]]) |
| tokens, mask, diag = comp(latents, frame_indices, None, target) |
| assert tokens.shape == (2, 2, 2, 8) |
| assert mask.shape == (2, 2, 2) |
| assert mask[0, 0].any() |
| assert diag["selected_source_count"].max().item() <= 4 |
|
|
|
|
| def test_dynamic_compressor_abstains_without_old_enough_sources(): |
| comp = small_compressor(exclude_latest_local_frames=4) |
| latents = torch.randn(2, 1, 3, 2, 2) |
| frame_indices = torch.tensor([[5], [6]]) |
| target = torch.tensor([[8]]) |
| tokens, mask, diag = comp(latents, frame_indices, None, target) |
| assert tokens.shape == (1, 1, 2, 8) |
| assert not mask.any() |
| assert diag["max_source_frame"].item() == -1 |
| assert diag["dynamic_min_gap_to_target_per_target"].item() == -1 |
|
|
|
|
| def test_dynamic_compressor_reports_generated_fraction_and_no_future(): |
| comp = small_compressor(exclude_latest_local_frames=0) |
| latents = torch.randn(3, 1, 3, 2, 2) |
| frame_indices = torch.tensor([[0], [2], [5]]) |
| generated = torch.tensor([[False], [True], [True]]) |
| target = torch.tensor([[3]]) |
| _, mask, diag = comp(latents, frame_indices, None, target, generated) |
| assert mask.any() |
| assert diag["max_source_frame"].item() == 2 |
| assert 0.0 < diag["generated_source_fraction"].item() < 1.0 |
|
|
|
|
| def test_dynamic_compressor_excludes_c_short_overlap_and_keeps_shape(): |
| comp = small_compressor(exclude_latest_local_frames=2) |
| latents = torch.randn(5, 1, 3, 2, 2) |
| frame_indices = torch.tensor([[0], [1], [2], [3], [4]]) |
| target = torch.tensor([[5]]) |
| tokens, mask, diag = comp(latents, frame_indices, None, target) |
| assert tokens.shape == (1, 1, 2, 8) |
| assert mask.any() |
| assert diag["max_source_frame"].item() == 2 |
| assert diag["dynamic_min_gap_to_target_per_target"].item() == 3 |
| assert diag["dynamic_max_gap_to_target_per_target"].item() == 5 |
| assert diag["dynamic_exclude_latest_local_frames"] == 2 |
|
|
|
|
| def test_cache_materialize_raw_latents_excludes_c_short_overlap(): |
| cache = StreamingCache(enabled=True, keep_raw_latents="all", keep_compressed_records=False) |
| latents = torch.randn(6, 1, 3, 2, 2) |
| frames = torch.arange(6).view(6, 1) |
| cache.add_raw_latents(latents, frames) |
| raw_latents, raw_frames, raw_generated, raw_pose = cache.materialize_raw_latents( |
| device=torch.device("cpu"), |
| dtype=latents.dtype, |
| max_recent_frames=8, |
| target_frame_indices=torch.tensor([[6]]), |
| exclude_latest_local_frames=4, |
| ) |
| assert raw_pose is None |
| assert raw_latents.shape[0] == 2 |
| assert raw_generated.shape == raw_frames.shape |
| assert raw_frames.flatten().tolist() == [0, 1] |
|
|
|
|
| def test_dynamic_compressor_preserves_grad_to_trainable_parts(): |
| comp = small_compressor(exclude_latest_local_frames=0) |
| latents = torch.randn(4, 1, 3, 2, 2) |
| frame_indices = torch.arange(4)[:, None] |
| target = torch.tensor([[4]]) |
| tokens, mask, _ = comp(latents, frame_indices, None, target) |
| assert mask.any() |
| tokens.square().sum().backward() |
| grads = [ |
| comp.conv3d.weight.grad, |
| comp.out_norm.weight.grad, |
| ] |
| assert all(grad is not None for grad in grads) |
| assert all(grad.abs().sum().item() > 0 for grad in grads) |
|
|
|
|
| def test_dynamic_compressor_selects_only_recent_valid_sources(): |
| comp = small_compressor(exclude_latest_local_frames=2) |
| latents = torch.randn(20, 1, 3, 2, 2) |
| frame_indices = torch.arange(20)[:, None] |
| target = torch.tensor([[10]]) |
| _, mask, diag = comp(latents, frame_indices, None, target) |
| assert mask.any() |
| assert diag["selected_source_count"].item() == 4 |
|
|
