common

class ParallelTiledVAE(ABC):
    tile_sample_min_height: int
    tile_sample_min_width: int
    tile_sample_min_num_frames: int
    tile_sample_stride_height: int
    tile_sample_stride_width: int
    tile_sample_stride_num_frames: int
    blend_num_frames: int
    use_tiling: bool
    use_temporal_tiling: bool
    use_parallel_tiling: bool
def __init__(self, config: VAEConfig, **kwargs) -> None:
    self.config = config
    self.tile_sample_min_height = config.tile_sample_min_height
    self.tile_sample_min_width = config.tile_sample_min_width
    self.tile_sample_min_num_frames = config.tile_sample_min_num_frames
    self.tile_sample_stride_height = config.tile_sample_stride_height
    self.tile_sample_stride_width = config.tile_sample_stride_width
    self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames
    self.blend_num_frames = config.blend_num_frames
    self.use_tiling = config.use_tiling
    self.use_temporal_tiling = config.use_temporal_tiling
    self.use_parallel_tiling = config.use_parallel_tiling

def to(self, device) -> 'ParallelTiledVAE':
    return self

@property
def temporal_compression_ratio(self) -> int:
    return cast(int, self.config.temporal_compression_ratio)

@property
def spatial_compression_ratio(self) -> int:
    return cast(int, self.config.spatial_compression_ratio)

@property
def scaling_factor(self) -> float | torch.Tensor:
    return cast(float | torch.Tensor, self.config.scaling_factor)

@abstractmethod
def _encode(self, *args, **kwargs) -> torch.Tensor:
    pass

@abstractmethod
def _decode(self, *args, **kwargs) -> torch.Tensor:
    pass

def encode(self, x: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = x.shape
    latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1

    if self.use_tiling and self.use_temporal_tiling and num_frames > self.tile_sample_min_num_frames:
        latents = self.tiled_encode(x)[:, :, :latent_num_frames]
    elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
        latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames]
    else:
        latents = self._encode(x)[:, :, :latent_num_frames]
    return DiagonalGaussianDistribution(latents)

def decode(self, z: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = z.shape
    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
    num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1

    if self.use_tiling and self.use_parallel_tiling and get_sp_world_size() > 1:
        return self.parallel_tiled_decode(z)[:, :, :num_sample_frames]
    if self.use_tiling and self.use_temporal_tiling and num_frames > tile_latent_min_num_frames:
        return self.tiled_decode(z)[:, :, :num_sample_frames]

    if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
        return self.spatial_tiled_decode(z)[:, :, :num_sample_frames]

    return self._decode(z)[:, :, :num_sample_frames]

def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
    for y in range(blend_extent):
        b[:, :, :,
          y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent)
    return b

def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
    for x in range(blend_extent):
        b[:, :, :, :,
          x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent)
    return b

def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
    for x in range(blend_extent):
        b[:, :,
          x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :,
                                                                                   x, :, :] * (x / blend_extent)
    return b

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)

def _parallel_data_generator(self, gathered_results, gathered_dim_metadata) -> Iterator[tuple[torch.Tensor, int]]:
    global_idx = 0
    for i, per_rank_metadata in enumerate(gathered_dim_metadata):
        _start_shape = 0
        for shape in per_rank_metadata:
            mul_shape = prod(shape)
            yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx)
            _start_shape += mul_shape
            global_idx += 1

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 _merge_spatial_tiles(self, tiles, blend_height, blend_width, stride_height, stride_width) -> torch.Tensor:
    """Helper function to merge spatial tiles with blending"""
    result_rows = []
    for i, row in enumerate(tiles):
        result_row = []
        for j, tile in enumerate(row):
            if i > 0:
                tile = self.blend_v(tiles[i - 1][j], tile, blend_height)
            if j > 0:
                tile = self.blend_h(row[j - 1], tile, blend_width)
            result_row.append(tile[:, :, :, :stride_height, :stride_width])
        result_rows.append(torch.cat(result_row, dim=-1))
    return torch.cat(result_rows, dim=-2)

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 tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
    _, _, num_frames, height, width = x.shape

    # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
    tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio

    row = []
    for i in range(0, num_frames, self.tile_sample_stride_num_frames):
        tile = x[:, :, i:i + self.tile_sample_min_num_frames + 1, :, :]
        if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width):
            tile = self.spatial_tiled_encode(tile)
        else:
            tile = self._encode(tile)
        if i > 0:
            tile = tile[:, :, 1:, :, :]
        row.append(tile)
    blend_latent_num_frames = max(
        1,
        self.blend_num_frames // self.temporal_compression_ratio,
    )
    result_row = []
    for i, tile in enumerate(row):
        if i > 0:
            tile = self.blend_t(row[i - 1], tile, blend_latent_num_frames)
            result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
        else:
            result_row.append(tile[:, :, :tile_latent_stride_num_frames + 1, :, :])
    enc = torch.cat(result_row, dim=2)
    return enc

def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = z.shape

    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_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio

    row = []
    for i in range(0, num_frames, tile_latent_stride_num_frames):
        tile = z[:, :, i:i + tile_latent_min_num_frames + 1, :, :]
        if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height):
            decoded = self.spatial_tiled_decode(tile)
        else:
            decoded = self._decode(tile)
        if i > 0:
            decoded = decoded[:, :, 1:, :, :]
        row.append(decoded)
    result_row = []
    for i, tile in enumerate(row):
        if i > 0:
            tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
            result_row.append(tile[:, :, :self.tile_sample_stride_num_frames, :, :])
        else:
            result_row.append(tile[:, :, :self.tile_sample_stride_num_frames + 1, :, :])

    dec = torch.cat(result_row, dim=2)
    return dec

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 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</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.to
to(device) -> ParallelTiledVAE
Source
def to(self, device) -> 'ParallelTiledVAE':
    return self
fastvideo.models.vaes.common.ParallelTiledVAE.encode
encode(x: torch.Tensor) -> torch.Tensor
Source
def encode(self, x: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = x.shape
    latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
if self.use_tiling and self.use_temporal_tiling and num_frames &gt; self.tile_sample_min_num_frames:
    latents = self.tiled_encode(x)&#91;:, :, :latent_num_frames&#93;
elif self.use_tiling and (width &gt; self.tile_sample_min_width or height &gt; self.tile_sample_min_height):
    latents = self.spatial_tiled_encode(x)&#91;:, :, :latent_num_frames&#93;
else:
    latents = self._encode(x)&#91;:, :, :latent_num_frames&#93;
return DiagonalGaussianDistribution(latents)</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.decode
decode(z: torch.Tensor) -> torch.Tensor
Source
def decode(self, z: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = z.shape
    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
    num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
if self.use_tiling and self.use_parallel_tiling and get_sp_world_size() &gt; 1:
    return self.parallel_tiled_decode(z)&#91;:, :, :num_sample_frames&#93;
if self.use_tiling and self.use_temporal_tiling and num_frames &gt; tile_latent_min_num_frames:
    return self.tiled_decode(z)&#91;:, :, :num_sample_frames&#93;

if self.use_tiling and (width &gt; tile_latent_min_width or height &gt; tile_latent_min_height):
    return self.spatial_tiled_decode(z)&#91;:, :, :num_sample_frames&#93;

return self._decode(z)&#91;:, :, :num_sample_frames&#93;</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.blend_v
blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor
Source
def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
    for y in range(blend_extent):
        b[:, :, :,
          y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent)
    return b
fastvideo.models.vaes.common.ParallelTiledVAE.blend_h
blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor
Source
def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
    for x in range(blend_extent):
        b[:, :, :, :,
          x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent)
    return b
fastvideo.models.vaes.common.ParallelTiledVAE.blend_t
blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor
Source
def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
    blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
    for x in range(blend_extent):
        b[:, :,
          x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :,
                                                                                   x, :, :] * (x / blend_extent)
    return b
fastvideo.models.vaes.common.ParallelTiledVAE.spatial_tiled_encode
spatial_tiled_encode(x: torch.Tensor) -> torch.Tensor
Encode a batch of images using a tiled encoder.
Parameters

x: `torch.Tensor`
Input batch of videos.

Returns

torch.Tensor
torch.Tensor:
The latent representation of the encoded videos.

Source
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.
&quot;&quot;&quot;
_, _, _, 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 = &#91;&#93;
for i in range(0, height, self.tile_sample_stride_height):
    row = &#91;&#93;
    for j in range(0, width, self.tile_sample_stride_width):
        tile = x&#91;:, :, :, i:i + self.tile_sample_min_height, j:j + self.tile_sample_min_width&#93;
        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)</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.parallel_tiled_decode
parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensor
Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
Source
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 = &#91;&#93;
local_dim_metadata = &#91;&#93;
# 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&#91;:, :, 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&#93;

    # Process tile
    tile = self._decode(tile)

    if t_start &gt; 0:
        tile = tile&#91;:, :, 1:, :, :&#93;

    # 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(&#91;results.size(0)&#93;, device=results.device, dtype=torch.int64)
all_sizes = &#91;torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)&#93;
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&#91;:results.size(0)&#93; = results
del results

# Gather all results
gathered_dim_metadata = &#91;None&#93; * world_size
gathered_results = torch.zeros_like(padded_results).repeat(
    world_size, *&#91;1&#93; * len(padded_results.shape)).contiguous(
    )  # use contiguous to make sure it won&#x27;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 = &#91;&#91;&#91;&#91;&#93; for _ in range(num_w_tiles)&#93; for _ in range(num_h_tiles)&#93; for _ in range(num_t_tiles)&#93;
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&#91;t_idx&#93;&#91;h_idx&#93;&#91;w_idx&#93; = current_data
# Merge results
result_slices = &#91;&#93;
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 &gt; 0:
        slice_data = self.blend_t(last_slice_data, slice_data, self.blend_num_frames)
        result_slices.append(slice_data&#91;:, :, :self.tile_sample_stride_num_frames, :, :&#93;)
    else:
        result_slices.append(slice_data&#91;:, :, :self.tile_sample_stride_num_frames + 1, :, :&#93;)
    last_slice_data = slice_data
dec = torch.cat(result_slices, dim=2)

return dec</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.spatial_tiled_decode
spatial_tiled_decode(z: torch.Tensor) -> torch.Tensor
Decode a batch of images using a tiled decoder.
Parameters

z: `torch.Tensor`
Input batch of latent vectors.

Returns

torch.Tensor
torch.Tensor:
The decoded images.

Source
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.
&quot;&quot;&quot;

_, _, _, 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 = &#91;&#93;
for i in range(0, height, tile_latent_stride_height):
    row = &#91;&#93;
    for j in range(0, width, tile_latent_stride_width):
        tile = z&#91;:, :, :, i:i + tile_latent_min_height, j:j + tile_latent_min_width&#93;
        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)</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.tiled_encode
tiled_encode(x: torch.Tensor) -> torch.Tensor
Source
def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
    _, _, num_frames, height, width = x.shape
# tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio

row = &#91;&#93;
for i in range(0, num_frames, self.tile_sample_stride_num_frames):
    tile = x&#91;:, :, i:i + self.tile_sample_min_num_frames + 1, :, :&#93;
    if self.use_tiling and (height &gt; self.tile_sample_min_height or width &gt; self.tile_sample_min_width):
        tile = self.spatial_tiled_encode(tile)
    else:
        tile = self._encode(tile)
    if i &gt; 0:
        tile = tile&#91;:, :, 1:, :, :&#93;
    row.append(tile)
blend_latent_num_frames = max(
    1,
    self.blend_num_frames // self.temporal_compression_ratio,
)
result_row = &#91;&#93;
for i, tile in enumerate(row):
    if i &gt; 0:
        tile = self.blend_t(row&#91;i - 1&#93;, tile, blend_latent_num_frames)
        result_row.append(tile&#91;:, :, :tile_latent_stride_num_frames, :, :&#93;)
    else:
        result_row.append(tile&#91;:, :, :tile_latent_stride_num_frames + 1, :, :&#93;)
enc = torch.cat(result_row, dim=2)
return enc</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.tiled_decode
tiled_decode(z: torch.Tensor) -> torch.Tensor
Source
def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = z.shape
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_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio

row = &#91;&#93;
for i in range(0, num_frames, tile_latent_stride_num_frames):
    tile = z&#91;:, :, i:i + tile_latent_min_num_frames + 1, :, :&#93;
    if self.use_tiling and (tile.shape&#91;-1&#93; &gt; tile_latent_min_width or tile.shape&#91;-2&#93; &gt; tile_latent_min_height):
        decoded = self.spatial_tiled_decode(tile)
    else:
        decoded = self._decode(tile)
    if i &gt; 0:
        decoded = decoded&#91;:, :, 1:, :, :&#93;
    row.append(decoded)
result_row = &#91;&#93;
for i, tile in enumerate(row):
    if i &gt; 0:
        tile = self.blend_t(row&#91;i - 1&#93;, tile, self.blend_num_frames)
        result_row.append(tile&#91;:, :, :self.tile_sample_stride_num_frames, :, :&#93;)
    else:
        result_row.append(tile&#91;:, :, :self.tile_sample_stride_num_frames + 1, :, :&#93;)

dec = torch.cat(result_row, dim=2)
return dec</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.enable_tiling
enable_tiling(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
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.
Parameters

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.

Source
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.
&quot;&quot;&quot;
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</code></pre></details>
fastvideo.models.vaes.common.ParallelTiledVAE.disable_tiling
disable_tiling() -> None
Disable tiled VAE decoding. If enable_tiling was previously enabled, this method will go back to computing
decoding in one step.
Source
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
fastvideo.models.vaes.common.DiagonalGaussianDistribution
DiagonalGaussianDistribution(parameters: torch.Tensor, deterministic: bool = False)
Source
class DiagonalGaussianDistribution:
def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
    self.parameters = parameters
    self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
    self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
    self.deterministic = deterministic
    self.std = torch.exp(0.5 * self.logvar)
    self.var = torch.exp(self.logvar)
    if self.deterministic:
        self.var = self.std = torch.zeros_like(self.mean,
                                               device=self.parameters.device,
                                               dtype=self.parameters.dtype)

def sample(self, generator: torch.Generator | None = None) -&gt; torch.Tensor:
    # make sure sample is on the same device as the parameters and has same dtype
    sample = randn_tensor(
        self.mean.shape,
        generator=generator,
        device=self.parameters.device,
        dtype=self.parameters.dtype,
    )
    x = self.mean + self.std * sample
    return x

def kl(self, other: Optional&#91;&quot;DiagonalGaussianDistribution&quot;&#93; = None) -&gt; torch.Tensor:
    if self.deterministic:
        return torch.Tensor(&#91;0.0&#93;)
    else:
        if other is None:
            return 0.5 * torch.sum(
                torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
                dim=&#91;1, 2, 3&#93;,
            )
        else:
            return 0.5 * torch.sum(
                torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar +
                other.logvar,
                dim=&#91;1, 2, 3&#93;,
            )

def nll(self, sample: torch.Tensor, dims: tuple&#91;int, ...&#93; = (1, 2, 3)) -&gt; torch.Tensor:
    if self.deterministic:
        return torch.Tensor(&#91;0.0&#93;)
    logtwopi = np.log(2.0 * np.pi)
    return 0.5 * torch.sum(
        logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
        dim=dims,
    )

def mode(self) -&gt; torch.Tensor:
    return self.mean</code></pre></details>
fastvideo.models.vaes.common.DiagonalGaussianDistribution.sample
sample(generator: torch.Generator | None = None) -> torch.Tensor
Source
def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
    # make sure sample is on the same device as the parameters and has same dtype
    sample = randn_tensor(
        self.mean.shape,
        generator=generator,
        device=self.parameters.device,
        dtype=self.parameters.dtype,
    )
    x = self.mean + self.std * sample
    return x
fastvideo.models.vaes.common.DiagonalGaussianDistribution.kl
kl(other: Optional[DiagonalGaussianDistribution] = None) -> torch.Tensor
Source
def kl(self, other: Optional["DiagonalGaussianDistribution"] = None) -> torch.Tensor:
    if self.deterministic:
        return torch.Tensor([0.0])
    else:
        if other is None:
            return 0.5 * torch.sum(
                torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
                dim=[1, 2, 3],
            )
        else:
            return 0.5 * torch.sum(
                torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar +
                other.logvar,
                dim=[1, 2, 3],
            )
fastvideo.models.vaes.common.DiagonalGaussianDistribution.nll
nll(sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor
Source
def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
    if self.deterministic:
        return torch.Tensor([0.0])
    logtwopi = np.log(2.0 * np.pi)
    return 0.5 * torch.sum(
        logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
        dim=dims,
    )
fastvideo.models.vaes.common.DiagonalGaussianDistribution.mode
mode() -> torch.Tensor
Source
def mode(self) -> torch.Tensor:
    return self.mean