| from abc import ABC, abstractmethod |
| from typing import Tuple |
|
|
| import torch |
| from einops import rearrange |
| from torch import Tensor |
|
|
|
|
| def latent_to_pixel_coords( |
| latent_coords: Tensor, scale_factors: Tuple[int, int, int], causal_fix: bool = False |
| ) -> Tensor: |
| """ |
| Converts latent coordinates to pixel coordinates by scaling them according to the VAE's |
| configuration. |
| Args: |
| latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents] |
| containing the latent corner coordinates of each token. |
| scale_factors (Tuple[int, int, int]): The scale factors of the VAE's latent space. |
| causal_fix (bool): Whether to take into account the different temporal scale |
| of the first frame. Default = False for backwards compatibility. |
| Returns: |
| Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates. |
| """ |
| pixel_coords = ( |
| latent_coords |
| * torch.tensor(scale_factors, device=latent_coords.device)[None, :, None] |
| ) |
| if causal_fix: |
| |
| pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0) |
| return pixel_coords |
|
|
|
|
| class Patchifier(ABC): |
| def __init__(self, patch_size: int): |
| super().__init__() |
| self._patch_size = (1, patch_size, patch_size) |
|
|
| @abstractmethod |
| def patchify( |
| self, latents: Tensor, frame_rates: Tensor, scale_grid: bool |
| ) -> Tuple[Tensor, Tensor]: |
| pass |
|
|
| @abstractmethod |
| def unpatchify( |
| self, |
| latents: Tensor, |
| output_height: int, |
| output_width: int, |
| output_num_frames: int, |
| out_channels: int, |
| ) -> Tuple[Tensor, Tensor]: |
| pass |
|
|
| @property |
| def patch_size(self): |
| return self._patch_size |
|
|
| def get_latent_coords( |
| self, latent_num_frames, latent_height, latent_width, batch_size, device |
| ): |
| """ |
| Return a tensor of shape [batch_size, 3, num_patches] containing the |
| top-left corner latent coordinates of each latent patch. |
| The tensor is repeated for each batch element. |
| """ |
| latent_sample_coords = torch.meshgrid( |
| torch.arange(0, latent_num_frames, self._patch_size[0], device=device), |
| torch.arange(0, latent_height, self._patch_size[1], device=device), |
| torch.arange(0, latent_width, self._patch_size[2], device=device), |
| indexing="ij", |
| ) |
| latent_sample_coords = torch.stack(latent_sample_coords, dim=0) |
| latent_coords = latent_sample_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1) |
| latent_coords = rearrange( |
| latent_coords, "b c f h w -> b c (f h w)", b=batch_size |
| ) |
| return latent_coords |
|
|
|
|
| class SymmetricPatchifier(Patchifier): |
| def patchify( |
| self, |
| latents: Tensor, |
| ) -> Tuple[Tensor, Tensor]: |
| b, _, f, h, w = latents.shape |
| latent_coords = self.get_latent_coords(f, h, w, b, latents.device) |
| latents = rearrange( |
| latents, |
| "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", |
| p1=self._patch_size[0], |
| p2=self._patch_size[1], |
| p3=self._patch_size[2], |
| ) |
| return latents, latent_coords |
|
|
| def unpatchify( |
| self, |
| latents: Tensor, |
| output_height: int, |
| output_width: int, |
| output_num_frames: int, |
| out_channels: int, |
| ) -> Tuple[Tensor, Tensor]: |
| output_height = output_height // self._patch_size[1] |
| output_width = output_width // self._patch_size[2] |
| latents = rearrange( |
| latents, |
| "b (f h w) (c p q) -> b c f (h p) (w q) ", |
| f=output_num_frames, |
| h=output_height, |
| w=output_width, |
| p=self._patch_size[1], |
| q=self._patch_size[2], |
| ) |
| return latents |
|
|
