| 319 | self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() |
| 320 | |
| 321 | def window_attention(self, x, attn_layer, index): |
| 322 | attn_mask = None |
| 323 | if self.window_size is not None: |
| 324 | l = x.shape[1] |
| 325 | assert l % self.window_size == 0, "Sequence length must be divisible by window size" |
| 326 | if index % 2 == 0: |
| 327 | # Even index: split into windows without shifting. |
| 328 | x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size) |
| 329 | x = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 330 | x = rearrange(x, '(b p) w c -> b (p w) c', p=l // self.window_size) |
| 331 | else: |
| 332 | # Odd index: roll by half a window, then split. |
| 333 | x = torch.roll(x, shifts=self.window_size // 2, dims=1) |
| 334 | x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size) |
| 335 | x = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 336 | x = rearrange(x, '(b p) w c -> b (p w) c', p=l // self.window_size) |
| 337 | x = torch.roll(x, shifts=-self.window_size // 2, dims=1) |
| 338 | else: |
| 339 | x = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0] |
| 340 | return x |
| 341 | |
| 342 | def forward(self, x, index, condition): |
| 343 | residual = x.type(torch.float32) |