hugging-apps/echo-memory

1import torch2 3 4def apply_framepack_token_weights(5    x: torch.Tensor,6    num_context_frames: int,7    f: int,8    h: int,9    w: int,10    context_position: str = "prefix",11    use_framepack_memory: bool = False,12    context_temporal_decay: float = 1.0,13    context_attention_weight: float = 1.0,14):15    if x is None or x.ndim != 3:16        return x17    if not use_framepack_memory or int(num_context_frames) <= 0:18        return x19    b, n, d = x.shape20    f = int(f)21    if f <= 0 or n != f * int(h) * int(w):22        return x23 24    hw = int(h) * int(w)25    x4 = x.reshape(b, f, hw, d)26    k = min(int(num_context_frames), f)27    decay = float(context_temporal_decay)28    gain = float(context_attention_weight)29    if context_position == "suffix":30        ctx_start = f - k31        ctx_end = f32        # Suffix: first context frame is nearest boundary to target.33        distances = torch.arange(k, device=x.device, dtype=x.dtype)34    else:35        ctx_start = 036        ctx_end = k37        # Prefix: last context frame is nearest boundary to target.38        distances = torch.arange(k - 1, -1, -1, device=x.device, dtype=x.dtype)39 40    weights = gain * torch.pow(torch.tensor(decay, device=x.device, dtype=x.dtype), distances)41    x4[:, ctx_start:ctx_end, :, :] = x4[:, ctx_start:ctx_end, :, :] * weights.view(1, k, 1, 1)42    return x4.reshape(b, n, d)43 44