| import torch |
| from dememwm_import_helper import install_dememwm_namespace |
|
|
| install_dememwm_namespace() |
| from algorithms.worldmem.dememwm.algorithm import MemoryDiTMixin |
| from algorithms.worldmem.dememwm.cache import StreamingCache |
| from algorithms.worldmem.dememwm.types import MemoryRecord, MemorySourceType |
|
|
|
|
| class Harness(MemoryDiTMixin): |
| def __init__(self): |
| self.n_tokens = 8 |
| self.context_frames = 0 |
| self.frame_stack = 1 |
| self.dememwm_anchor_proj = torch.nn.Linear(12, 8) |
| self.dememwm_revisit_proj = torch.nn.Linear(12, 8) |
| self.project_call_lengths = [] |
| self.project_call_values = [] |
|
|
| def _project_latent_patch_tokens(self, latents, projection, patch_size): |
| self.project_call_lengths.append(int(latents.shape[0])) |
| self.project_call_values.append(latents[:, 0, 0, 0, 0].detach().cpu().tolist()) |
| return MemoryDiTMixin._project_latent_patch_tokens(self, latents, projection, patch_size) |
|
|
|
|
| def test_training_window_bounds_samples_inside_long_clip(): |
| harness = Harness() |
| torch.manual_seed(0) |
| starts = [] |
| for _ in range(20): |
| start, end = harness._training_window_bounds(128, torch.device("cpu")) |
| starts.append(start) |
| assert end - start == 8 |
| assert 0 <= start <= 120 |
| assert any(start != 120 for start in starts) |
|
|
|
|
| def test_training_window_bounds_respects_context_frames(): |
| harness = Harness() |
| harness.context_frames = 100 |
| torch.manual_seed(0) |
| starts = [] |
| for _ in range(20): |
| start, end = harness._training_window_bounds(128, torch.device("cpu")) |
| starts.append(start) |
| assert end - start == 8 |
| assert 100 <= start <= 120 |
| assert any(start != 120 for start in starts) |
|
|
|
|
| def test_revisit_local_context_exclusion_uses_n_tokens_times_frame_stack(): |
| harness = Harness() |
| harness.n_tokens = 4 |
| harness.frame_stack = 2 |
| harness.context_frames = 100 |
| assert harness._local_context_exclusion_frames() == 8 |
|
|
|
|
| def test_diverse_anchor_selection_does_not_repeat_tied_pose_indices(): |
| harness = Harness() |
| source_positions = torch.arange(5) |
| poses = torch.zeros((5, 5), dtype=torch.float32) |
|
|
| selected = harness._select_diverse_anchor_positions(source_positions, poses, 4) |
|
|
| assert selected.tolist() == [0, 1, 2, 3] |
|
|
|
|
| def test_diverse_anchor_selection_seeds_from_widest_pose_pair(): |
| harness = Harness() |
| source_positions = torch.arange(4) |
| poses = torch.tensor([[0.0], [-10.0], [10.0], [0.1]], dtype=torch.float32) |
|
|
| selected = harness._select_diverse_anchor_positions(source_positions, poses, 2) |
|
|
| assert selected.tolist() == [1, 2] |
|
|
|
|
| def test_cached_revisit_prefilter_keeps_only_causal_records(): |
| harness = Harness() |
|
|
| def record(frame: int) -> MemoryRecord: |
| return MemoryRecord( |
| tokens=torch.zeros((1, 8)), |
| mask=torch.ones(1, dtype=torch.bool), |
| source_start=frame, |
| source_end=frame + 1, |
| frame_indices=torch.tensor([frame]), |
| pose=None, |
| source_type=MemorySourceType.REVISIT, |
| is_generated=False, |
| chunk_id=f"revisit_{frame}", |
| ) |
|
|
| selected = harness._causal_cached_revisit_records( |
| (record(0), record(2), record(5)), |
| target_frame=3, |
| ) |
|
|
| assert [record.source_start for record in selected] == [0, 2] |
|
|
|
|
| def test_diverse_anchor_selection_uses_context_frames_not_literal_limit(): |
| harness = Harness() |
| harness.context_frames = 2 |
| latents = torch.randn(8, 1, 3, 2, 2) |
| frame_indices = torch.arange(8)[:, None] |
| poses = torch.zeros((8, 1, 5), dtype=torch.float32) |
| target_pose = torch.zeros((1, 1, 5), dtype=torch.float32) |
| anchor_banks, _, _, diag = harness._build_preselected_causal_memory_banks( |
| committed_latents=latents, |
| source_frame_indices=frame_indices, |
| source_is_generated=None, |
| pose=poses, |
| action=None, |
| target_frame_indices=torch.tensor([[6]]), |
| target_pose=target_pose, |
| target_action=None, |
| target_video_ids=None, |
| allow_generated_anchor=False, |
| anchor_indices=[0, 1, 2, 3], |
| anchor_pool_h=1, |
| anchor_pool_w=1, |
| anchor_diverse=True, |
| revisit_pool_h=1, |
| revisit_pool_w=1, |
| revisit_max_frames=0, |
| exclude_local_context_frames=4, |
| fov_overlap_threshold=0.0, |
| plucker_weight=0.1, |
| revisit_retrieval_kwargs=None, |
| token_patch_size=2, |
| ) |
|
|
| assert [int(record.frame_indices.item()) for record in anchor_banks[0].records] == [0, 1] |
| assert diag["preselected_anchor_projected_frame_count"] == 2 |
|
|
|
|
| def test_streaming_diverse_anchor_selection_uses_context_frames(): |
| harness = Harness() |
| harness.context_frames = 2 |
| latents = torch.randn(8, 1, 3, 2, 2) |
| frame_indices = torch.arange(8)[:, None] |
| poses = torch.zeros((8, 1, 5), dtype=torch.float32) |
|
|
| anchor_banks, _ = harness._build_streaming_cache_records( |
| source_latents=latents, |
| source_frame_indices=frame_indices, |
| source_is_generated=None, |
| pose=poses, |
| action=None, |
| allow_generated_anchor=False, |
| anchor_indices=[0, 1, 2, 3], |
| anchor_pool_h=1, |
| anchor_pool_w=1, |
| anchor_diverse=True, |
| token_patch_size=2, |
| ) |
|
|
| assert [int(record.frame_indices.item()) for record in anchor_banks[0].records] == [0, 1] |
| assert harness.project_call_lengths == [2] |
|
|
|
|
| def test_preselected_memory_banks_project_only_selected_frames(): |
| harness = Harness() |
| latents = torch.randn(20, 1, 3, 2, 2) |
| frame_indices = torch.arange(20)[:, None] |
| target_frame_indices = torch.tensor([[10], [11]]) |
| poses = torch.zeros((20, 1, 5), dtype=torch.float32) |
| target_pose = torch.zeros((2, 1, 5), dtype=torch.float32) |
| anchor_banks, revisit_banks, tokens_per_frame, diag = harness._build_preselected_causal_memory_banks( |
| committed_latents=latents, |
| source_frame_indices=frame_indices, |
| source_is_generated=None, |
| pose=poses, |
| action=None, |
| target_frame_indices=target_frame_indices, |
| target_pose=target_pose, |
| target_action=None, |
| target_video_ids=None, |
| allow_generated_anchor=False, |
| anchor_indices=[0, 1, 2, 3], |
| anchor_pool_h=1, |
| anchor_pool_w=1, |
| anchor_diverse=False, |
| revisit_pool_h=1, |
| revisit_pool_w=1, |
| revisit_max_frames=2, |
| exclude_local_context_frames=4, |
| fov_overlap_threshold=0.0, |
| plucker_weight=0.1, |
| revisit_retrieval_kwargs=None, |
| token_patch_size=2, |
| ) |
| assert tokens_per_frame == 1 |
| assert len(anchor_banks[0].records) == 4 |
| assert len(revisit_banks[0].records) == 3 |
| assert diag["preselected_anchor_projected_frame_count"] == 4 |
| assert diag["preselected_revisit_projected_frame_count"] == 3 |
| assert diag["preselected_revisit_projected_frame_record_count"] == 3 |
| assert harness.project_call_lengths == [4, 1, 1, 1] |
| assert 20 not in harness.project_call_lengths |
|
|
|
|
| def test_preselected_revisit_projects_best_fov_frame_not_latest(): |
| harness = Harness() |
| latents = torch.arange(8, dtype=torch.float32).view(8, 1, 1, 1, 1).expand(8, 1, 3, 2, 2).clone() |
| frame_indices = torch.arange(8)[:, None] |
| pose_rows = torch.tensor( |
| [ |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 0.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| [0.0, 0.0, 0.0, 0.0, 180.0], |
| ], |
| dtype=torch.float32, |
| ) |
| poses = pose_rows[:, None, :] |
|
|
| _, revisit_banks, _, _ = harness._build_preselected_causal_memory_banks( |
| committed_latents=latents, |
| source_frame_indices=frame_indices, |
| source_is_generated=None, |
| pose=poses, |
| action=None, |
| target_frame_indices=torch.tensor([[8]]), |
| target_pose=torch.tensor([[[0.0, 0.0, 0.0, 0.0, 0.0]]]), |
| target_action=None, |
| target_video_ids=None, |
| allow_generated_anchor=False, |
| anchor_indices=[], |
| anchor_pool_h=1, |
| anchor_pool_w=1, |
| anchor_diverse=False, |
| revisit_pool_h=1, |
| revisit_pool_w=1, |
| revisit_max_frames=1, |
| exclude_local_context_frames=4, |
| fov_overlap_threshold=0.30, |
| plucker_weight=0.1, |
| revisit_retrieval_kwargs={"high_quality_fov_threshold": 0.70}, |
| token_patch_size=2, |
| ) |
|
|
| assert len(revisit_banks[0].records) == 1 |
| assert revisit_banks[0].records[0].metadata["dememwm_selected_frame_index"] == 1 |
| assert harness.project_call_values == [[1.0]] |
|
|
|
|
| def test_streaming_revisit_projection_uses_selected_frame_metadata(): |
| harness = Harness() |
| cache = StreamingCache(enabled=True, keep_raw_latents="all", keep_compressed_records=True) |
| latents = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 1, 3, 2, 2).clone() |
| cache.add_raw_latents(latents, torch.arange(4)[:, None]) |
| record = MemoryRecord( |
| tokens=torch.zeros((1, 8)), |
| mask=torch.ones(1, dtype=torch.bool), |
| source_start=0, |
| source_end=4, |
| frame_indices=torch.tensor([0, 1, 2, 3]), |
| pose=None, |
| source_type=MemorySourceType.PREFIX_GT, |
| is_generated=False, |
| chunk_id="frame", |
| metadata={ |
| "dememwm_revisit_metadata_only": True, |
| "dememwm_selected_frame_index": 1, |
| }, |
| ) |
|
|
| projected = harness._project_streaming_revisit_records( |
| cache=cache, |
| batch_idx=0, |
| records=[record], |
| device=torch.device("cpu"), |
| dtype=torch.float32, |
| token_patch_size=2, |
| revisit_pool_h=1, |
| revisit_pool_w=1, |
| projection_cache={}, |
| ) |
|
|
| assert len(projected) == 1 |
| assert projected[0].metadata["dememwm_selected_frame_index"] == 1 |
| assert harness.project_call_values == [[1.0]] |
|
|
