common ¶

def disable_tiling(self) -> None:
    r"""
    Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
    decoding in one step.
    """
    self.use_tiling = False

def enable_tiling(
    self,
    tile_sample_min_height: int | None = None,
    tile_sample_min_width: int | None = None,
    tile_sample_min_num_frames: int | None = None,
    tile_sample_stride_height: int | None = None,
    tile_sample_stride_width: int | None = None,
    tile_sample_stride_num_frames: int | None = None,
    blend_num_frames: int | None = None,
    use_tiling: bool | None = None,
    use_temporal_tiling: bool | None = None,
    use_parallel_tiling: bool | None = None,
) -> None:
    r"""
    Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
    compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
    processing larger images.

    Args:
        tile_sample_min_height (`int`, *optional*):
            The minimum height required for a sample to be separated into tiles across the height dimension.
        tile_sample_min_width (`int`, *optional*):
            The minimum width required for a sample to be separated into tiles across the width dimension.
        tile_sample_min_num_frames (`int`, *optional*):
            The minimum number of frames required for a sample to be separated into tiles across the frame
            dimension.
        tile_sample_stride_height (`int`, *optional*):
            The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
            no tiling artifacts produced across the height dimension.
        tile_sample_stride_width (`int`, *optional*):
            The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
            artifacts produced across the width dimension.
        tile_sample_stride_num_frames (`int`, *optional*):
            The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts
            produced across the frame dimension.
    """
    self.use_tiling = True
    self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
    self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
    self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames
    self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
    self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
    self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
    if blend_num_frames is not None:
        self.blend_num_frames = blend_num_frames
    else:
        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
    self.use_tiling = use_tiling or self.use_tiling
    self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling
    self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling

def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
    """
    Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
    """
    world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
    B, C, T, H, W = z.shape

    # Calculate parameters
    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
    tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
    tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio

    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width

    # Calculate tile dimensions
    num_t_tiles = (T + tile_latent_stride_num_frames -
                   1) // tile_latent_stride_num_frames
    num_h_tiles = (H + tile_latent_stride_height -
                   1) // tile_latent_stride_height
    num_w_tiles = (W + tile_latent_stride_width -
                   1) // tile_latent_stride_width
    total_spatial_tiles = num_h_tiles * num_w_tiles
    total_tiles = num_t_tiles * total_spatial_tiles

    # Calculate tiles per rank and padding
    tiles_per_rank = (total_tiles + world_size - 1) // world_size
    start_tile_idx = rank * tiles_per_rank
    end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)

    local_results = []
    local_dim_metadata = []
    # Process assigned tiles
    for local_idx, global_idx in enumerate(
            range(start_tile_idx, end_tile_idx)):
        t_idx = global_idx // total_spatial_tiles
        spatial_idx = global_idx % total_spatial_tiles
        h_idx = spatial_idx // num_w_tiles
        w_idx = spatial_idx % num_w_tiles

        # Calculate positions
        t_start = t_idx * tile_latent_stride_num_frames
        h_start = h_idx * tile_latent_stride_height
        w_start = w_idx * tile_latent_stride_width

        # Extract and process tile
        tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1,
                 h_start:h_start + tile_latent_min_height,
                 w_start:w_start + tile_latent_min_width]

        # Process tile
        tile = self._decode(tile)

        if t_start > 0:
            tile = tile[:, :, 1:, :, :]

        # Store metadata
        shape = tile.shape
        # Store decoded data (flattened)
        decoded_flat = tile.reshape(-1)
        local_results.append(decoded_flat)
        local_dim_metadata.append(shape)

    results = torch.cat(local_results, dim=0).contiguous()
    del local_results
    # first gather size to pad the results
    local_size = torch.tensor([results.size(0)],
                              device=results.device,
                              dtype=torch.int64)
    all_sizes = [
        torch.zeros(1, device=results.device, dtype=torch.int64)
        for _ in range(world_size)
    ]
    dist.all_gather(all_sizes, local_size)
    max_size = max(size.item() for size in all_sizes)
    padded_results = torch.zeros(max_size, device=results.device)
    padded_results[:results.size(0)] = results
    del results

    # Gather all results
    gathered_dim_metadata = [None] * world_size
    gathered_results = torch.zeros_like(padded_results).repeat(
        world_size, *[1] * len(padded_results.shape)
    ).contiguous(
    )  # use contiguous to make sure it won't copy data in the following operations
    # TODO (PY): use fastvideo distributed methods
    dist.all_gather_into_tensor(gathered_results, padded_results)
    dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
    # Process gathered results
    data: list = [[[[] for _ in range(num_w_tiles)]
                   for _ in range(num_h_tiles)] for _ in range(num_t_tiles)]
    for current_data, global_idx in self._parallel_data_generator(
            gathered_results, gathered_dim_metadata):
        t_idx = global_idx // total_spatial_tiles
        spatial_idx = global_idx % total_spatial_tiles
        h_idx = spatial_idx // num_w_tiles
        w_idx = spatial_idx % num_w_tiles
        data[t_idx][h_idx][w_idx] = current_data
    # Merge results
    result_slices = []
    last_slice_data = None
    for i, tem_data in enumerate(data):
        slice_data = self._merge_spatial_tiles(
            tem_data, blend_height, blend_width,
            self.tile_sample_stride_height, self.tile_sample_stride_width)
        if i > 0:
            slice_data = self.blend_t(last_slice_data, slice_data,
                                      self.blend_num_frames)
            result_slices.append(
                slice_data[:, :, :self.tile_sample_stride_num_frames, :, :])
        else:
            result_slices.append(
                slice_data[:, :, :self.tile_sample_stride_num_frames +
                           1, :, :])
        last_slice_data = slice_data
    dec = torch.cat(result_slices, dim=2)

    return dec

def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
    r"""
    Decode a batch of images using a tiled decoder.

    Args:
        z (`torch.Tensor`): Input batch of latent vectors.

    Returns:
        `torch.Tensor`:
            The decoded images.
    """

    _, _, _, height, width = z.shape
    # sample_height = height * self.spatial_compression_ratio
    # sample_width = width * self.spatial_compression_ratio

    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio

    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width

    # Split z into overlapping tiles and decode them separately.
    # The tiles have an overlap to avoid seams between tiles.
    rows = []
    for i in range(0, height, tile_latent_stride_height):
        row = []
        for j in range(0, width, tile_latent_stride_width):
            tile = z[:, :, :, i:i + tile_latent_min_height,
                     j:j + tile_latent_min_width]
            decoded = self._decode(tile)
            row.append(decoded)
        rows.append(row)
    return self._merge_spatial_tiles(rows, blend_height, blend_width,
                                     self.tile_sample_stride_height,
                                     self.tile_sample_stride_width)

def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
    r"""Encode a batch of images using a tiled encoder.

    Args:
        x (`torch.Tensor`): Input batch of videos.

    Returns:
        `torch.Tensor`:
            The latent representation of the encoded videos.
    """
    _, _, _, height, width = x.shape
    # latent_height = height // self.spatial_compression_ratio
    # latent_width = width // self.spatial_compression_ratio

    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio

    blend_height = tile_latent_min_height - tile_latent_stride_height
    blend_width = tile_latent_min_width - tile_latent_stride_width

    # Split x into overlapping tiles and encode them separately.
    # The tiles have an overlap to avoid seams between tiles.
    rows = []
    for i in range(0, height, self.tile_sample_stride_height):
        row = []
        for j in range(0, width, self.tile_sample_stride_width):
            tile = x[:, :, :, i:i + self.tile_sample_min_height,
                     j:j + self.tile_sample_min_width]
            tile = self._encode(tile)
            row.append(tile)
        rows.append(row)

    return self._merge_spatial_tiles(rows, blend_height, blend_width,
                                     tile_latent_stride_height,
                                     tile_latent_stride_width)