Method forward

wan/modules/s2v/motioner.py:189–239

(self, x, seq_lens, grid_sizes, freqs)

class SwinSelfAttention(SelfAttention):

    def forward(self, x, seq_lens, grid_sizes, freqs):
        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
        assert b == 1, 'Only support batch_size 1'

        # query, key, value function
        def qkv_fn(x):
            q = self.norm_q(self.q(x)).view(b, s, n, d)
            k = self.norm_k(self.k(x)).view(b, s, n, d)
            v = self.v(x).view(b, s, n, d)
            return q, k, v

        q, k, v = qkv_fn(x)

        q = rope_apply(q, grid_sizes, freqs)
        k = rope_apply(k, grid_sizes, freqs)
        T, H, W = grid_sizes[0].tolist()

        q = rearrange(q, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
        k = rearrange(k, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
        v = rearrange(v, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)

        ref_q = q[-1:]
        q = q[:-1]

        ref_k = repeat(
            k[-1:], "1 s n d -> t s n d", t=k.shape[0] - 1)  # t hw n d
        k = k[:-1]
        k = torch.cat([k[:1], k, k[-1:]])
        k = torch.cat([k[1:-1], k[2:], k[:-2], ref_k], dim=1)  # (bt) (3hw) n d

        ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=v.shape[0] - 1)
        v = v[:-1]
        v = torch.cat([v[:1], v, v[-1:]])
        v = torch.cat([v[1:-1], v[2:], v[:-2], ref_v], dim=1)

        # q: b (t h w) n d
        # k: b (t h w) n d
        out = flash_attention(
            q=q,
            k=k,
            v=v,
            # k_lens=torch.tensor([k.shape[1]] * k.shape[0], device=x.device, dtype=torch.long),
            window_size=self.window_size)
        out = torch.cat([out, ref_v[:1]], axis=0)
        out = rearrange(out, '(b t) (h w) n d -> b (t h w) n d', t=T, h=H, w=W)
        x = out

        # output
        x = x.flatten(2)
        x = self.o(x)
        return x


#Fix the reference frame RoPE to 1,H,W.