File size: 2,448 Bytes
2ea2835 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import copy
from types import SimpleNamespace
import pytest
from model.wan_kv_mmd import WanKVMMD
def _triple_state_only():
model = object.__new__(WanKVMMD)
model.feature_type = "triple_dinov3_vjepa2_videomae_v2_dense_latent"
model.is_dinov3_dense = False
model.is_triple_dense = True
model.uses_fresh_video_gradcache = True
model.dense_feature_prefix = "triple"
model.dense_encoder_name = "triple"
model.dense_log_prefix = "triple_dense"
model.vjepa2_multilayer_enabled = False
model.mmd_estimator = "nystrom_exact_repulsion"
model.triple_teacher_bank_content_hash = "a" * 64
model.dual_teacher_bank_content_hash = "a" * 64
model.dual_branch_registry = {
"dinov3": {
"reference_path": "/data/dino.pt",
"reference_hash": "b" * 64,
"nystrom_landmarks": 256,
},
"vjepa2": {
"reference_path": "/data/vj.pt",
"reference_hash": "c" * 64,
"nystrom_landmarks": 1024,
},
"videomae_v2": {
"reference_path": "/data/vm.pt",
"reference_hash": "d" * 64,
"nystrom_landmarks": 256,
},
}
model.latent_mmd = SimpleNamespace(
attraction_ema_enabled=False,
repulsion_fifo_enabled=False,
)
model.distribution_state_pending = lambda: False
model._fd_metadata = lambda: {
"feature_type": model.feature_type,
"teacher_bank_content_hash": model.triple_teacher_bank_content_hash,
}
return model
def test_triple_fd_state_round_trip_uses_own_namespace_and_version_10():
model = _triple_state_only()
state = model.fd_state_dict()
assert state["triple_dense_nystrom_mmd_version"] == 10
assert "vjepa2_dense_nystrom_mmd_version" not in state
model.load_fd_state_dict(copy.deepcopy(state))
assert model.can_load_initialization_fd_state(state)
def test_triple_fd_state_rejects_foreign_or_wrong_version():
model = _triple_state_only()
state = model.fd_state_dict()
wrong = copy.deepcopy(state)
wrong["triple_dense_nystrom_mmd_version"] = 9
with pytest.raises(ValueError, match="triple_dense_nystrom_mmd_version"):
model.load_fd_state_dict(wrong)
foreign = copy.deepcopy(state)
foreign["vjepa2_dense_nystrom_mmd_version"] = 10
with pytest.raises(ValueError, match="another dense feature namespace"):
model.load_fd_state_dict(foreign)
|