File size: 1,523 Bytes
c335050 | 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 | """FramePack-Length latent/RT alignment (e.g. K=5, r=4). Run: PYTHONPATH=. python3 tests/test_framepack_memory_align.py"""
import torch
from diffsynth.models.memory.framepack_length import (
framepack_align_context_actions_to_latents,
framepack_length_compress_context_latents,
)
from diffsynth.models.memory.framepack_weight import apply_framepack_token_weights
def test_k5_r4_latent_and_actions():
B, C, H, W = 1, 16, 8, 8
K = 5
r = 4
lat = torch.randn(B, C, K, H, W)
out, new_k, K_pad, K_orig = framepack_length_compress_context_latents(lat, r)
assert K_orig == 5
pad = (r - (K % r)) % r
assert K_pad == K + pad == 8
assert new_k == 2
assert out.shape[2] == 2
ca = torch.randn(K, 12)
aligned = framepack_align_context_actions_to_latents(
ca, K_orig, K_pad, r, device=lat.device, dtype=lat.dtype
)
assert aligned.shape == (2, 12)
def test_framepack_weight_preserves_shape_suffix():
D = 64
f, h, w = 5, 2, 2
num_ctx = 2
N = f * h * w
x = torch.randn(1, N, D)
y = apply_framepack_token_weights(
x,
num_context_frames=num_ctx,
f=f,
h=h,
w=w,
context_position="suffix",
use_framepack_memory=True,
context_temporal_decay=0.9,
context_attention_weight=1.0,
)
assert y.shape == x.shape
if __name__ == "__main__":
test_k5_r4_latent_and_actions()
test_framepack_weight_preserves_shape_suffix()
print("test_framepack_memory_align: ok")
|