Add DeMemWM frame memory reference attention

algorithms/dememwm/models/dit.py

@@ -8,6 +8,7 @@ References:
 from typing import Optional, Literal
 import torch
 from torch import nn
+from torch.nn import functional as F
 from .rotary_embedding_torch import RotaryEmbedding
 from einops import rearrange
 from .attention import SpatialAxialAttention, TemporalAxialAttention, MemTemporalAxialAttention, MemFullAttention
@@ -140,6 +141,64 @@ class FinalLayer(nn.Module):
         return x
 
 
+class FrameMemoryReferenceAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.to_q = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.to_k = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.to_v = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def _split_heads(self, x):
+        return rearrange(x, "b n (h d) -> b h n d", h=self.num_heads)
+
+    def forward(self, target_hidden, memory_hidden, memory_mask=None):
+        B, T_target, H, W, D = target_hidden.shape
+        if memory_hidden is None or int(memory_hidden.shape[1]) == 0:
+            return target_hidden.new_zeros(target_hidden.shape)
+
+        T_memory = memory_hidden.shape[1]
+        P = H * W
+        q = rearrange(target_hidden, "b t h w d -> (b t) (h w) d")
+        q = self._split_heads(self.to_q(q))
+
+        memory_tokens = rearrange(memory_hidden, "b m h w d -> b (m h w) d")
+        k = self._split_heads(self.to_k(memory_tokens))
+        v = self._split_heads(self.to_v(memory_tokens))
+
+        # A zero dummy key keeps all-padded samples finite; the projected delta
+        # is masked back to zero below so padded memory cannot contribute.
+        dummy = k.new_zeros((B, self.num_heads, 1, self.head_dim))
+        k = torch.cat([k, dummy], dim=2)
+        v = torch.cat([v, dummy], dim=2)
+
+        if memory_mask is None:
+            frame_valid = torch.ones((B, T_memory), device=target_hidden.device, dtype=torch.bool)
+        else:
+            frame_valid = memory_mask.to(device=target_hidden.device, dtype=torch.bool)
+        key_valid = frame_valid[:, :, None].expand(B, T_memory, P).reshape(B, T_memory * P)
+        active = key_valid.any(dim=1)
+        key_valid = torch.cat([key_valid, ~active[:, None]], dim=1)
+
+        row_batch = torch.arange(B, device=target_hidden.device).repeat_interleave(T_target)
+        attn_bias = torch.zeros((B * T_target, 1, 1, key_valid.shape[1]), device=q.device, dtype=q.dtype)
+        attn_bias = attn_bias.masked_fill(~key_valid[row_batch][:, None, None], float("-inf"))
+
+        x = F.scaled_dot_product_attention(
+            query=q.contiguous(),
+            key=k.index_select(0, row_batch).contiguous(),
+            value=v.index_select(0, row_batch).contiguous(),
+            attn_mask=attn_bias,
+        )
+        x = rearrange(x, "n h p d -> n p (h d)")
+        x = self.out_proj(x).to(dtype=target_hidden.dtype)
+        x = rearrange(x, "(b t) (h w) d -> b t h w d", b=B, t=T_target, h=H, w=W)
+        return x * active[:, None, None, None, None].to(dtype=x.dtype)
+
+
 class SpatioTemporalDiTBlock(nn.Module):
     def __init__(
         self,
@@ -240,9 +299,70 @@ class SpatioTemporalDiTBlock(nn.Module):
         if self.ref_mode == 'parallel':
             self.parallel_map = nn.Linear(hidden_size, hidden_size)
 
+        if self.use_memory_attention:
+            self.r_attn_anchor = FrameMemoryReferenceAttention(hidden_size, num_heads)
+            self.r_attn_dynamic = FrameMemoryReferenceAttention(hidden_size, num_heads)
+            self.r_attn_revisit = FrameMemoryReferenceAttention(hidden_size, num_heads)
+
+    def _split_frame_memory(self, x, frame_memory_segments):
+        target = int(frame_memory_segments.get("target", 0))
+        anchor = int(frame_memory_segments.get("anchor", 0))
+        dynamic = int(frame_memory_segments.get("dynamic", 0))
+        revisit = int(frame_memory_segments.get("revisit", 0))
+        a0, a1 = target, target + anchor
+        d0, d1 = a1, a1 + dynamic
+        r0, r1 = d1, d1 + revisit
+        return x[:, :target], x[:, a0:a1], x[:, d0:d1], x[:, r0:r1]
+
+    def _frame_memory_stream_mask(self, frame_memory_masks, stream_name, stream_hidden):
+        if stream_hidden is None or int(stream_hidden.shape[1]) == 0:
+            return None
+        if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None:
+            return None
+        return frame_memory_masks[stream_name].to(device=stream_hidden.device, dtype=torch.bool)
+
+    def _apply_frame_memory_reference_attention(self, x, c, frame_memory_segments, frame_memory_masks):
+        x_target, x_anchor, x_dynamic, x_revisit = self._split_frame_memory(x, frame_memory_segments)
+        if int(x_target.shape[1]) == 0:
+            return x
+
+        r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
+        attn_input = modulate(self.r_norm1(x), r_shift_msa, r_scale_msa)
+        attn_target, attn_anchor, attn_dynamic, attn_revisit = self._split_frame_memory(attn_input, frame_memory_segments)
+
+        deltas = []
+        active_stream_count = x_target.new_zeros((x_target.shape[0],))
+        for stream_name, r_attn, stream_hidden, stream_attn in (
+            ("anchor", self.r_attn_anchor, x_anchor, attn_anchor),
+            ("dynamic", self.r_attn_dynamic, x_dynamic, attn_dynamic),
+            ("revisit", self.r_attn_revisit, x_revisit, attn_revisit),
+        ):
+            if int(stream_hidden.shape[1]) == 0:
+                continue
+            stream_mask = self._frame_memory_stream_mask(frame_memory_masks, stream_name, stream_hidden)
+            deltas.append(r_attn(attn_target, stream_attn, stream_mask))
+            if stream_mask is None:
+                active_stream_count = active_stream_count + 1
+            else:
+                active_stream_count = active_stream_count + stream_mask.any(dim=1).to(dtype=active_stream_count.dtype)
+
+        if deltas:
+            # Average over streams that have at least one valid memory frame per
+            # sample; samples with no active streams receive a zero reference delta.
+            r_delta = sum(deltas)
+            active = active_stream_count > 0
+            r_delta = r_delta / active_stream_count.clamp_min(1)[:, None, None, None, None]
+            r_delta = r_delta * active[:, None, None, None, None].to(dtype=r_delta.dtype)
+            x_target = x_target + gate(r_delta, r_gate_msa[:, :x_target.shape[1]])
+
+        x = torch.cat([x_target, x_anchor, x_dynamic, x_revisit], dim=1)
+        x = x + gate(self.r_mlp(modulate(self.r_norm2(x), r_shift_mlp, r_scale_mlp)), r_gate_mlp)
+        return x
+
     def forward(self, x, c, current_frame=None, timestep=None, is_last_block=False, 
         pose_cond=None, mode="training", c_action_cond=None, reference_length=None,
-        frame_memory_segments=None, frame_memory_masks=None):
+        frame_memory_segments=None, frame_memory_masks=None, frame_memory_pose=None,
+        image_hw=None):
         B, T, H, W, D = x.shape
 
         # spatial block
@@ -270,7 +390,14 @@ class SpatioTemporalDiTBlock(nn.Module):
         # memory block
         relative_embedding = self.relative_embedding # and mode == "training"
 
-        if self.use_memory_attention:
+        if self.use_memory_attention and frame_memory_segments is not None:
+            x = self._apply_frame_memory_reference_attention(
+                x,
+                c,
+                frame_memory_segments,
+                frame_memory_masks,
+            )
+        elif self.use_memory_attention:
             r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
 
             if pose_cond is not None:
@@ -547,7 +674,9 @@ class DiT(nn.Module):
             x = block(x, c, current_frame=current_frame, timestep=t, is_last_block= (i+1 == len(self.blocks)), 
                 pose_cond=pc, mode=mode, c_action_cond=c_action_cond, reference_length=reference_length,
                 frame_memory_segments=frame_memory_segments,
-                frame_memory_masks=frame_memory_masks)  # (N, T, H, W, D)
+                frame_memory_masks=frame_memory_masks,
+                frame_memory_pose=frame_memory_pose,
+                image_hw=image_hw)  # (N, T, H, W, D)
         x = self.final_layer(x, c) # (N, T, H, W, patch_size ** 2 * out_channels)
         # unpatchify
         x = rearrange(x, "b t h w d -> (b t) h w d")

tests/test_dememwm_temporal_attention.py

@@ -4,7 +4,7 @@ import torch
 from torch import nn
 
 from algorithms.dememwm.models.attention import TemporalAxialAttention
-from algorithms.dememwm.models.dit import DiT, SpatioTemporalDiTBlock
+from algorithms.dememwm.models.dit import DiT, FrameMemoryReferenceAttention, SpatioTemporalDiTBlock
 
 
 class IdentityRotary:
@@ -134,6 +134,25 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertIs(spy.calls[0]["frame_memory_segments"], segments)
         self.assertIs(spy.calls[0]["frame_memory_masks"], masks)
 
+    def test_frame_memory_reference_attention_masks_padded_keys(self):
+        attn = FrameMemoryReferenceAttention(hidden_size=1, num_heads=1)
+        with torch.no_grad():
+            attn.to_q.weight.zero_()
+            attn.to_k.weight.zero_()
+            attn.to_v.weight.fill_(1.0)
+            attn.out_proj.weight.fill_(1.0)
+            attn.out_proj.bias.zero_()
+
+        target = torch.zeros((1, 1, 1, 1, 1))
+        memory = torch.tensor([1.0, 10.0]).view(1, 2, 1, 1, 1)
+
+        first_only = attn(target, memory, torch.tensor([[True, False]]))
+        second_only = attn(target, memory, torch.tensor([[False, True]]))
+        empty = attn(target, memory, torch.zeros((1, 2), dtype=torch.bool))
+
+        self.assertTrue(torch.allclose(first_only, torch.ones_like(first_only)))
+        self.assertTrue(torch.allclose(second_only, torch.full_like(second_only, 10.0)))
+        self.assertTrue(torch.equal(empty, torch.zeros_like(empty)))
 
 if __name__ == "__main__":
     unittest.main()