x: (N, T/2 * S, patch_size**3 * C) imgs: (N, T, H, W, C)
(x, c, p, w, h, rope_position_ids=None, **kwargs)
| 440 | |
| 441 | |
| 442 | def unpatchify(x, c, p, w, h, rope_position_ids=None, **kwargs): |
| 443 | """ |
| 444 | x: (N, T/2 * S, patch_size**3 * C) |
| 445 | imgs: (N, T, H, W, C) |
| 446 | """ |
| 447 | if rope_position_ids is not None: |
| 448 | assert NotImplementedError |
| 449 | # do pix2struct unpatchify |
| 450 | L = x.shape[1] |
| 451 | x = x.reshape(shape=(x.shape[0], L, p, p, c)) |
| 452 | x = torch.einsum("nlpqc->ncplq", x) |
| 453 | imgs = x.reshape(shape=(x.shape[0], c, p, L * p)) |
| 454 | else: |
| 455 | b = x.shape[0] |
| 456 | imgs = rearrange(x, "b (t h w) (c p q) -> b t c (h p) (w q)", b=b, h=h, w=w, c=c, p=p, q=p) |
| 457 | |
| 458 | return imgs |
| 459 | |
| 460 | |
| 461 | class FinalLayerMixin(BaseMixin): |