r""" Reconstruct video tensors from patch embeddings. Args: x (List[Tensor]): List of patchified features, each with shape [L, C_out * prod(patch_size)] grid_sizes (Tensor): Original spatial-temporal grid dimensions before patc
(self, x, grid_sizes)
| 492 | # Removed forward() because we don't use it due to pipeline parallelism. |
| 493 | |
| 494 | def unpatchify(self, x, grid_sizes): |
| 495 | r""" |
| 496 | Reconstruct video tensors from patch embeddings. |
| 497 | |
| 498 | Args: |
| 499 | x (List[Tensor]): |
| 500 | List of patchified features, each with shape [L, C_out * prod(patch_size)] |
| 501 | grid_sizes (Tensor): |
| 502 | Original spatial-temporal grid dimensions before patching, |
| 503 | shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) |
| 504 | |
| 505 | Returns: |
| 506 | List[Tensor]: |
| 507 | Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] |
| 508 | """ |
| 509 | |
| 510 | c = self.out_dim |
| 511 | out = [] |
| 512 | for u, v in zip(x, grid_sizes.tolist()): |
| 513 | u = u[:math.prod(v)].view(*v, *self.patch_size, c) |
| 514 | u = torch.einsum('fhwpqrc->cfphqwr', u) |
| 515 | u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) |
| 516 | out.append(u) |
| 517 | return out |
| 518 | |
| 519 | def init_weights(self): |
| 520 | r""" |