| import unittest |
| from unittest.mock import patch |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from predictor_training.atc_fusion import ATCFusion, MultiInputMLP, StateTransport |
| from predictor_training.single_block import SingleBlockPredictor |
| from wan.modules.model import rope_params |
|
|
|
|
| def small_freqs(head_dim: int = 8) -> torch.Tensor: |
| return torch.cat( |
| [ |
| rope_params(32, head_dim - 4 * (head_dim // 6)), |
| rope_params(32, 2 * (head_dim // 6)), |
| rope_params(32, 2 * (head_dim // 6)), |
| ], |
| dim=1, |
| ) |
|
|
|
|
| def reference_attention( |
| q: torch.Tensor, k: torch.Tensor, v: torch.Tensor |
| ) -> torch.Tensor: |
| scores = torch.einsum("blhd,bshd->bhls", q, k) / q.shape[-1] ** 0.5 |
| probability = scores.float().softmax(dim=-1).to(v.dtype) |
| return torch.einsum("bhls,bshd->blhd", probability, v) |
|
|
|
|
| class DummyBlock(nn.Module): |
| num_heads = 1 |
|
|
| def forward(self, value: torch.Tensor, **_kwargs) -> torch.Tensor: |
| return value |
|
|
|
|
| class ATCFusionTest(unittest.TestCase): |
| def setUp(self) -> None: |
| torch.manual_seed(0) |
| self.grid = torch.tensor([[2, 2, 2]], dtype=torch.long) |
| self.freqs = small_freqs() |
| self.current = torch.randn(1, 8, 8) |
| self.previous = torch.randn(1, 8, 8) |
|
|
| def test_multi_input_mlp_matches_explicit_concat(self) -> None: |
| module = MultiInputMLP(3, dim=8, hidden_dim=12) |
| values = [torch.randn(2, 5, 8) for _ in range(3)] |
| combined_weight = torch.cat( |
| [layer.weight for layer in module.inputs], dim=1 |
| ) |
| combined_bias = sum(layer.bias for layer in module.inputs) |
| reference_hidden = F.linear( |
| torch.cat(values, dim=-1), combined_weight, combined_bias |
| ) |
| reference = module.out(module.activation(reference_hidden)) |
| torch.testing.assert_close(module(*values), reference) |
|
|
| @patch( |
| "predictor_training.atc_fusion.attention", |
| side_effect=reference_attention, |
| ) |
| def test_last_frame_scope_cannot_read_earlier_frames(self, _mock) -> None: |
| transport = StateTransport( |
| dim=8, |
| num_heads=1, |
| previous_scope="last_frame", |
| gradient_checkpointing=False, |
| diagnostic_queries=4, |
| ).eval() |
| first, _ = transport( |
| self.current, self.previous, self.grid, self.freqs, current_start=8 |
| ) |
| perturbed_early = self.previous.clone() |
| perturbed_early[:, :4] += torch.randn_like(perturbed_early[:, :4]) |
| second, _ = transport( |
| self.current, |
| perturbed_early, |
| self.grid, |
| self.freqs, |
| current_start=8, |
| ) |
| torch.testing.assert_close(first, second, rtol=0, atol=0) |
|
|
| perturbed_last = self.previous.clone() |
| perturbed_last[:, 4:] += torch.randn_like(perturbed_last[:, 4:]) |
| third, _ = transport( |
| self.current, |
| perturbed_last, |
| self.grid, |
| self.freqs, |
| current_start=8, |
| ) |
| self.assertFalse(torch.equal(first, third)) |
|
|
| @patch( |
| "predictor_training.atc_fusion.attention", |
| side_effect=reference_attention, |
| ) |
| def test_zero_initialized_correction_bypass(self, _mock) -> None: |
| fusion = ATCFusion( |
| dim=8, |
| num_heads=1, |
| freq_dim=8, |
| mlp_hidden_dim=16, |
| gate_hidden_dim=4, |
| previous_scope="chunk", |
| gradient_checkpointing=False, |
| ).eval() |
| condition = torch.randn_like(self.current) |
| backbone, gated_correction = fusion( |
| self.current, |
| torch.randn_like(self.current), |
| self.previous, |
| condition, |
| torch.tensor([62.5]), |
| self.grid, |
| self.freqs, |
| current_start=8, |
| ) |
| self.assertEqual(backbone.shape, self.current.shape) |
| torch.testing.assert_close( |
| gated_correction, |
| torch.zeros_like(gated_correction), |
| rtol=0, |
| atol=0, |
| ) |
| self.assertAlmostEqual(float(fusion.last_gate.mean()), 0.3, delta=0.01) |
|
|
| def test_atc_token_gate_is_a_stage1_parameter(self) -> None: |
| predictor = SingleBlockPredictor( |
| DummyBlock(), |
| dim=8, |
| gradient_checkpointing=False, |
| input_variant="atc", |
| atc_freq_dim=8, |
| atc_mlp_hidden_dim=16, |
| atc_gate_hidden_dim=4, |
| ) |
| token_gate_ids = { |
| id(parameter) |
| for parameter in predictor.fusion.token_gate.parameters() |
| } |
| stage1_ids = { |
| id(parameter) for parameter in predictor.stage1_input_parameters() |
| } |
| self.assertTrue(token_gate_ids) |
| self.assertTrue(token_gate_ids.issubset(stage1_ids)) |
| self.assertEqual(predictor.gate_parameters(), []) |
|
|
| @patch( |
| "predictor_training.atc_fusion.attention", |
| side_effect=reference_attention, |
| ) |
| def test_zero_initialized_predictor_returns_anchor(self, _mock) -> None: |
| predictor = SingleBlockPredictor( |
| DummyBlock(), |
| dim=8, |
| gradient_checkpointing=False, |
| input_variant="atc", |
| atc_freq_dim=8, |
| atc_mlp_hidden_dim=16, |
| atc_gate_hidden_dim=4, |
| ).eval() |
| anchor = torch.randn_like(self.current) |
| empty_history = torch.empty(1, 0, 1, 8) |
| prediction = predictor( |
| current_tokens=self.current, |
| anchor_hidden=anchor, |
| previous_hidden=self.previous, |
| timestep_modulation=torch.randn(1, 2, 6, 8), |
| grid_sizes=self.grid, |
| freqs=self.freqs, |
| history_k=empty_history, |
| history_v=empty_history, |
| cross_k=torch.empty(1, 0, 1, 8), |
| cross_v=torch.empty(1, 0, 1, 8), |
| current_start=8, |
| condition_tokens=torch.randn_like(self.current), |
| anchor_distance=torch.tensor([62.5]), |
| ) |
| torch.testing.assert_close(prediction, anchor, rtol=0, atol=0) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
