| 270 | self.window_size = window_size |
| 271 | |
| 272 | def attention(self, x, index): |
| 273 | attn_mask = None |
| 274 | if self.window_size is not None: |
| 275 | l = x.shape[1] |
| 276 | assert l % self.window_size == 0 |
| 277 | if index % 2 == 0: |
| 278 | x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size) |
| 279 | x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 280 | x = rearrange(x, '(b l) w c -> b (l w) c', l=l//self.window_size, w=self.window_size) |
| 281 | else: |
| 282 | x = torch.roll(x, shifts=self.window_size//2, dims=1) |
| 283 | x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size) |
| 284 | x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 285 | x = rearrange(x, '(b l) w c -> b (l w) c', l=l//self.window_size, w=self.window_size) |
| 286 | x = torch.roll(x, shifts=-self.window_size//2, dims=1) |
| 287 | else: |
| 288 | x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 289 | return x |
| 290 | |
| 291 | def forward(self, x, index, condition=None): |
| 292 | # no condition in encoder, its a dummy argument |