--- title: ltx2vae ---

LTX-2 Video VAE implementation

fastvideo.models.vaes.ltx2vae.DiagonalGaussianDistribution

DiagonalGaussianDistribution(parameters: torch.Tensor, deterministic: bool = False)

fastvideo.models.vaes.ltx2vae.DiagonalGaussianDistribution.sample

sample(generator: torch.Generator | None = None) -> torch.Tensor

fastvideo.models.vaes.ltx2vae.DiagonalGaussianDistribution.kl

kl(other: Optional[DiagonalGaussianDistribution] = None) -> torch.Tensor

fastvideo.models.vaes.ltx2vae.DiagonalGaussianDistribution.nll

nll(sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor

fastvideo.models.vaes.ltx2vae.DiagonalGaussianDistribution.mode

mode() -> torch.Tensor

fastvideo.models.vaes.ltx2vae.NormLayerType

NormLayerType()
Source
class NormLayerType(Enum):
    GROUP_NORM = "group_norm"
    PIXEL_NORM = "pixel_norm"

fastvideo.models.vaes.ltx2vae.LogVarianceType

LogVarianceType()
Source
class LogVarianceType(Enum):
    PER_CHANNEL = "per_channel"
    UNIFORM = "uniform"
    CONSTANT = "constant"
    NONE = "none"

fastvideo.models.vaes.ltx2vae.PaddingModeType

PaddingModeType()
Source
class PaddingModeType(Enum):
    ZEROS = "zeros"
    REFLECT = "reflect"
    REPLICATE = "replicate"
    CIRCULAR = "circular"

fastvideo.models.vaes.ltx2vae.SpatialTilingConfig

SpatialTilingConfig(tile_size_in_pixels: int, tile_overlap_in_pixels: int = 0)

Configuration for dividing each frame into spatial tiles with optional overlap.

Source
@dataclass(frozen=True)
class SpatialTilingConfig:
    """Configuration for dividing each frame into spatial tiles with optional overlap."""
    tile_size_in_pixels: int
    tile_overlap_in_pixels: int = 0

    def __post_init__(self) -> None:
        if self.tile_size_in_pixels < 64:
            raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}")
        if self.tile_size_in_pixels % 32 != 0:
            raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}")
        if self.tile_overlap_in_pixels % 32 != 0:
            raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}")
        if self.tile_overlap_in_pixels >= self.tile_size_in_pixels:
            raise ValueError(
                f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}"
            )

fastvideo.models.vaes.ltx2vae.TemporalTilingConfig

TemporalTilingConfig(tile_size_in_frames: int, tile_overlap_in_frames: int = 0)

Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.

Source
@dataclass(frozen=True)
class TemporalTilingConfig:
    """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap."""
    tile_size_in_frames: int
    tile_overlap_in_frames: int = 0

    def __post_init__(self) -> None:
        if self.tile_size_in_frames < 16:
            raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}")
        if self.tile_size_in_frames % 8 != 0:
            raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}")
        if self.tile_overlap_in_frames % 8 != 0:
            raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}")
        if self.tile_overlap_in_frames >= self.tile_size_in_frames:
            raise ValueError(
                f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}"
            )

fastvideo.models.vaes.ltx2vae.TilingConfig

TilingConfig(spatial_config: SpatialTilingConfig | None = None, temporal_config: TemporalTilingConfig | None = None)

Configuration for splitting video into tiles with optional overlap.

Source
@dataclass(frozen=True)
class TilingConfig:
    """Configuration for splitting video into tiles with optional overlap."""
    spatial_config: SpatialTilingConfig | None = None
    temporal_config: TemporalTilingConfig | None = None

    @classmethod
    def default(cls) -> "TilingConfig":
        return cls(
            spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
            temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
        )

fastvideo.models.vaes.ltx2vae.TilingConfig.default

default() -> 'TilingConfig'
Source
@classmethod
def default(cls) -> "TilingConfig":
    return cls(
        spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
        temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
    )

fastvideo.models.vaes.ltx2vae.DimensionIntervals

DimensionIntervals(starts: List[int], ends: List[int], left_ramps: List[int], right_ramps: List[int])

Intervals which a single dimension of the latent space is split into.

Source
@dataclass(frozen=True)
class DimensionIntervals:
    """Intervals which a single dimension of the latent space is split into."""
    starts: List[int]
    ends: List[int]
    left_ramps: List[int]
    right_ramps: List[int]

fastvideo.models.vaes.ltx2vae.LatentIntervals

LatentIntervals(original_shape: torch.Size, dimension_intervals: Tuple[DimensionIntervals, ...])

Intervals which the latent tensor of given shape is split into.

Source
@dataclass(frozen=True)
class LatentIntervals:
    """Intervals which the latent tensor of given shape is split into."""
    original_shape: torch.Size
    dimension_intervals: Tuple[DimensionIntervals, ...]

fastvideo.models.vaes.ltx2vae.VideoLatentShape

VideoLatentShape()

Shape of the tensor representing video in VAE latent space.

Source
class VideoLatentShape(NamedTuple):
    """Shape of the tensor representing video in VAE latent space."""
    batch: int
    channels: int
    frames: int
    height: int
    width: int

    def to_torch_shape(self) -> torch.Size:
        return torch.Size([self.batch, self.channels, self.frames, self.height, self.width])

    @staticmethod
    def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
        return VideoLatentShape(
            batch=shape[0],
            channels=shape[1],
            frames=shape[2],
            height=shape[3],
            width=shape[4],
        )

    def upscale(self, time_scale: int, spatial_scale: int) -> "VideoLatentShape":
        return self._replace(
            channels=3,
            frames=(self.frames - 1) * time_scale + 1,
            height=self.height * spatial_scale,
            width=self.width * spatial_scale,
        )

fastvideo.models.vaes.ltx2vae.VideoLatentShape.to_torch_shape

to_torch_shape() -> torch.Size
Source
def to_torch_shape(self) -> torch.Size:
    return torch.Size([self.batch, self.channels, self.frames, self.height, self.width])

fastvideo.models.vaes.ltx2vae.VideoLatentShape.from_torch_shape

from_torch_shape(shape: torch.Size) -> 'VideoLatentShape'
Source
@staticmethod
def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
    return VideoLatentShape(
        batch=shape[0],
        channels=shape[1],
        frames=shape[2],
        height=shape[3],
        width=shape[4],
    )

fastvideo.models.vaes.ltx2vae.VideoLatentShape.upscale

upscale(time_scale: int, spatial_scale: int) -> 'VideoLatentShape'
Source
def upscale(self, time_scale: int, spatial_scale: int) -> "VideoLatentShape":
    return self._replace(
        channels=3,
        frames=(self.frames - 1) * time_scale + 1,
        height=self.height * spatial_scale,
        width=self.width * spatial_scale,
    )

fastvideo.models.vaes.ltx2vae.Tile

Tile()

Represents a single tile.

Source
class Tile(NamedTuple):
    """Represents a single tile."""
    in_coords: Tuple[slice, ...]
    out_coords: Tuple[slice, ...]
    masks_1d: Tuple[torch.Tensor | None, ...]

    @property
    def blend_mask(self) -> torch.Tensor:
        num_dims = len(self.out_coords)
        per_dimension_masks: List[torch.Tensor] = []

        for dim_idx in range(num_dims):
            mask_1d = self.masks_1d[dim_idx]
            view_shape = [1] * num_dims
            if mask_1d is None:
                one = torch.ones(1)
                view_shape[dim_idx] = 1
                per_dimension_masks.append(one.view(*view_shape))
                continue

            view_shape[dim_idx] = mask_1d.shape[0]
            per_dimension_masks.append(mask_1d.view(*view_shape))

        combined_mask = per_dimension_masks[0]
        for mask in per_dimension_masks[1:]:
            combined_mask = combined_mask * mask

        return combined_mask

fastvideo.models.vaes.ltx2vae.compute_trapezoidal_mask_1d

compute_trapezoidal_mask_1d(length: int, ramp_left: int, ramp_right: int, left_starts_from_0: bool = False) -> torch.Tensor

Generate a 1D trapezoidal blending mask with linear ramps.

Source
def compute_trapezoidal_mask_1d(
    length: int,
    ramp_left: int,
    ramp_right: int,
    left_starts_from_0: bool = False,
) -> torch.Tensor:
    """Generate a 1D trapezoidal blending mask with linear ramps."""
    if length <= 0:
        raise ValueError("Mask length must be positive.")

    ramp_left = max(0, min(ramp_left, length))
    ramp_right = max(0, min(ramp_right, length))

    mask = torch.ones(length)

    if ramp_left > 0:
        interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
        fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
        if not left_starts_from_0:
            fade_in = fade_in[1:]
        mask[:ramp_left] *= fade_in

    if ramp_right > 0:
        fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
        mask[-ramp_right:] *= fade_out

    return mask.clamp_(0, 1)

fastvideo.models.vaes.ltx2vae.default_split_operation

default_split_operation(length: int) -> DimensionIntervals
Source
def default_split_operation(length: int) -> DimensionIntervals:
    return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])

fastvideo.models.vaes.ltx2vae.default_mapping_operation

default_mapping_operation(_intervals: DimensionIntervals) -> Tuple[List[slice], List[torch.Tensor | None]]
Source
def default_mapping_operation(_intervals: DimensionIntervals, ) -> Tuple[List[slice], List[torch.Tensor | None]]:
    return [slice(0, None)], [None]

fastvideo.models.vaes.ltx2vae.split_in_spatial

split_in_spatial(size: int, overlap: int) -> SplitOperation
Source
def split_in_spatial(size: int, overlap: int) -> SplitOperation:

    def split(dimension_size: int) -> DimensionIntervals:
        if dimension_size <= size:
            return DEFAULT_SPLIT_OPERATION(dimension_size)
        amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
        starts = [i * (size - overlap) for i in range(amount)]
        ends = [start + size for start in starts]
        ends[-1] = dimension_size
        left_ramps = [0] + [overlap] * (amount - 1)
        right_ramps = [overlap] * (amount - 1) + [0]
        return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)

    return split

fastvideo.models.vaes.ltx2vae.split_in_temporal

split_in_temporal(size: int, overlap: int) -> SplitOperation
Source
def split_in_temporal(size: int, overlap: int) -> SplitOperation:
    non_causal_split = split_in_spatial(size, overlap)

    def split(dimension_size: int) -> DimensionIntervals:
        if dimension_size <= size:
            return DEFAULT_SPLIT_OPERATION(dimension_size)
        intervals = non_causal_split(dimension_size)
        starts = list(intervals.starts)
        starts[1:] = [s - 1 for s in starts[1:]]
        left_ramps = list(intervals.left_ramps)
        left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
        return replace(intervals, starts=starts, left_ramps=left_ramps)

    return split

fastvideo.models.vaes.ltx2vae.map_temporal_slice

map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]
Source
def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
    start = begin * scale
    stop = 1 + (end - 1) * scale
    left_ramp_scaled = 1 + (left_ramp - 1) * scale
    right_ramp_scaled = right_ramp * scale

    return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp_scaled, right_ramp_scaled, True)

fastvideo.models.vaes.ltx2vae.map_spatial_slice

map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]
Source
def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
    start = begin * scale
    stop = end * scale
    left_ramp_scaled = left_ramp * scale
    right_ramp_scaled = right_ramp * scale

    return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp_scaled, right_ramp_scaled, False)

fastvideo.models.vaes.ltx2vae.to_mapping_operation

to_mapping_operation(map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]], scale: int) -> MappingOperation
Source
def to_mapping_operation(
    map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
    scale: int,
) -> MappingOperation:

    def map_op(intervals: DimensionIntervals) -> Tuple[List[slice], List[torch.Tensor | None]]:
        output_slices: List[slice] = []
        masks_1d: List[torch.Tensor | None] = []
        number_of_slices = len(intervals.starts)
        for i in range(number_of_slices):
            start = intervals.starts[i]
            end = intervals.ends[i]
            left_ramp = intervals.left_ramps[i]
            right_ramp = intervals.right_ramps[i]
            output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
            output_slices.append(output_slice)
            masks_1d.append(mask_1d)
        return output_slices, masks_1d

    return map_op

fastvideo.models.vaes.ltx2vae.create_tiles_from_intervals_and_mappers

create_tiles_from_intervals_and_mappers(intervals: LatentIntervals, mappers: List[MappingOperation]) -> List[Tile]
Source
def create_tiles_from_intervals_and_mappers(
    intervals: LatentIntervals,
    mappers: List[MappingOperation],
) -> List[Tile]:
    full_dim_input_slices = []
    full_dim_output_slices = []
    full_dim_masks_1d = []
    for axis_index in range(len(intervals.original_shape)):
        dimension_intervals = intervals.dimension_intervals[axis_index]
        starts = dimension_intervals.starts
        ends = dimension_intervals.ends
        input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
        output_slices, masks_1d = mappers[axis_index](dimension_intervals)
        full_dim_input_slices.append(input_slices)
        full_dim_output_slices.append(output_slices)
        full_dim_masks_1d.append(masks_1d)

    tiles = []
    tile_in_coords = list(itertools.product(*full_dim_input_slices))
    tile_out_coords = list(itertools.product(*full_dim_output_slices))
    tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
    for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
        tiles.append(Tile(
            in_coords=in_coord,
            out_coords=out_coord,
            masks_1d=mask_1d,
        ))
    return tiles

fastvideo.models.vaes.ltx2vae.create_tiles

create_tiles(latent_shape: torch.Size, splitters: List[SplitOperation], mappers: List[MappingOperation]) -> List[Tile]
Source
def create_tiles(
    latent_shape: torch.Size,
    splitters: List[SplitOperation],
    mappers: List[MappingOperation],
) -> List[Tile]:
    if len(splitters) != len(latent_shape):
        raise ValueError(f"Number of splitters must be equal to number of dimensions in latent shape, "
                         f"got {len(splitters)} and {len(latent_shape)}")
    if len(mappers) != len(latent_shape):
        raise ValueError(f"Number of mappers must be equal to number of dimensions in latent shape, "
                         f"got {len(mappers)} and {len(latent_shape)}")
    intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
    latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
    return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)

fastvideo.models.vaes.ltx2vae.patchify

patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor

Rearrange spatial dimensions into channels (space-to-depth).

Parameters

x: torch.Tensor

Input tensor (4D or 5D)

patch_size_hw: int

Spatial patch size for height and width.

patch_size_t: int

Temporal patch size. Default=1 (no temporal patching).

For 5D: (B, C, F, H, W) -> (B, Cpatch_size_hw^2patch_size_t, F/patch_size_t, H/patch_size_hw, W/patch_size_hw)

Source
def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
    """
    Rearrange spatial dimensions into channels (space-to-depth).

    Args:
        x: Input tensor (4D or 5D)
        patch_size_hw: Spatial patch size for height and width.
        patch_size_t: Temporal patch size. Default=1 (no temporal patching).

    For 5D: (B, C, F, H, W) -> (B, C*patch_size_hw^2*patch_size_t, F/patch_size_t, H/patch_size_hw, W/patch_size_hw)
    """
    if patch_size_hw == 1 and patch_size_t == 1:
        return x
    if x.dim() == 4:
        x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw)
    elif x.dim() == 5:
        x = rearrange(
            x,
            "b c (f p) (h q) (w r) -> b (c p r q) f h w",
            p=patch_size_t,
            q=patch_size_hw,
            r=patch_size_hw,
        )
    else:
        raise ValueError(f"Invalid input shape: {x.shape}")
    return x

fastvideo.models.vaes.ltx2vae.unpatchify

unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor

Rearrange channels back into spatial dimensions (depth-to-space).

Parameters

x: torch.Tensor

Input tensor (4D or 5D)

patch_size_hw: int

Spatial patch size for height and width.

patch_size_t: int

Temporal patch size. Default=1 (no temporal expansion).

For 5D: (B, Cpatch_size_hw^2patch_size_t, F, H, W) -> (B, C, Fpatch_size_t, Hpatch_size_hw, W*patch_size_hw)

Source
def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
    """
    Rearrange channels back into spatial dimensions (depth-to-space).

    Args:
        x: Input tensor (4D or 5D)
        patch_size_hw: Spatial patch size for height and width.
        patch_size_t: Temporal patch size. Default=1 (no temporal expansion).

    For 5D: (B, C*patch_size_hw^2*patch_size_t, F, H, W) -> (B, C, F*patch_size_t, H*patch_size_hw, W*patch_size_hw)
    """
    if patch_size_hw == 1 and patch_size_t == 1:
        return x

    if x.dim() == 4:
        x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw)
    elif x.dim() == 5:
        x = rearrange(
            x,
            "b (c p r q) f h w -> b c (f p) (h q) (w r)",
            p=patch_size_t,
            q=patch_size_hw,
            r=patch_size_hw,
        )
    return x

fastvideo.models.vaes.ltx2vae.PixelNorm

PixelNorm(dim: int = 1, eps: float = 1e-08)

Per-pixel (per-location) RMS normalization layer. Normalizes along the channel dimension using root-mean-square.

Source
class PixelNorm(nn.Module):
    """
    Per-pixel (per-location) RMS normalization layer.
    Normalizes along the channel dimension using root-mean-square.
    """

    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
        super().__init__()
        self.dim = dim
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
        return x / torch.sqrt(mean_sq + self.eps)

fastvideo.models.vaes.ltx2vae.PixelNorm.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
    return x / torch.sqrt(mean_sq + self.eps)

fastvideo.models.vaes.ltx2vae.PerChannelStatistics

PerChannelStatistics(latent_channels: int = 128)

Per-channel statistics for normalizing and denormalizing the latent representation. Statistics are computed over the dataset and stored in the model checkpoint.

Source
class PerChannelStatistics(nn.Module):
    """
    Per-channel statistics for normalizing and denormalizing the latent representation.
    Statistics are computed over the dataset and stored in the model checkpoint.
    """

    def __init__(self, latent_channels: int = 128):
        super().__init__()
        self.register_buffer("std-of-means", torch.ones(latent_channels))
        self.register_buffer("mean-of-means", torch.zeros(latent_channels))
        self.register_buffer("mean-of-stds", torch.ones(latent_channels))
        self.register_buffer("mean-of-stds_over_std-of-means", torch.ones(latent_channels))
        self.register_buffer("channel", torch.arange(latent_channels, dtype=torch.float32))

    def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
        std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
        mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
        return x * std + mean

    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
        mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
        return (x - mean) / std

fastvideo.models.vaes.ltx2vae.PerChannelStatistics.un_normalize

un_normalize(x: torch.Tensor) -> torch.Tensor
Source
def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
    std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
    mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
    return x * std + mean

fastvideo.models.vaes.ltx2vae.PerChannelStatistics.normalize

normalize(x: torch.Tensor) -> torch.Tensor
Source
def normalize(self, x: torch.Tensor) -> torch.Tensor:
    std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
    mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
    return (x - mean) / std

fastvideo.models.vaes.ltx2vae.CausalConv3d

CausalConv3d(in_channels: int, out_channels: int, kernel_size: int = 3, stride: int | Tuple[int, int, int] = 1, dilation: int = 1, groups: int = 1, bias: bool = True, spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS)

Causal 3D convolution that pads temporally by repeating the first frame.

Source
class CausalConv3d(nn.Module):
    """
    Causal 3D convolution that pads temporally by repeating the first frame.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int | Tuple[int, int, int] = 1,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = True,
        spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    ) -> None:
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        kernel_size = (kernel_size, kernel_size, kernel_size)
        self.time_kernel_size = kernel_size[0]

        dilation = (dilation, 1, 1)

        height_pad = kernel_size[1] // 2
        width_pad = kernel_size[2] // 2
        padding = (0, height_pad, width_pad)

        self.conv = nn.Conv3d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            dilation=dilation,
            padding=padding,
            padding_mode=spatial_padding_mode.value,
            groups=groups,
            bias=bias,
        )

    def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
        if causal:
            first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
            x = torch.cat((first_frame_pad, x), dim=2)
        else:
            first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
            last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
            x = torch.cat((first_frame_pad, x, last_frame_pad), dim=2)
        x = self.conv(x)
        return x

    @property
    def weight(self) -> torch.Tensor:
        return self.conv.weight

fastvideo.models.vaes.ltx2vae.CausalConv3d.forward

forward(x: torch.Tensor, causal: bool = True) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
    if causal:
        first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
        x = torch.cat((first_frame_pad, x), dim=2)
    else:
        first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
        last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
        x = torch.cat((first_frame_pad, x, last_frame_pad), dim=2)
    x = self.conv(x)
    return x

fastvideo.models.vaes.ltx2vae.make_conv_nd

make_conv_nd(dims: int, in_channels: int, out_channels: int, kernel_size: int, stride: int | Tuple[int, int, int] = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, causal: bool = False, spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS) -> nn.Module

Create a convolution layer (2D or 3D, causal or not).

Source
def make_conv_nd(
    dims: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride: int | Tuple[int, int, int] = 1,
    padding: int = 0,
    dilation: int = 1,
    groups: int = 1,
    bias: bool = True,
    causal: bool = False,
    spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
) -> nn.Module:
    """Create a convolution layer (2D or 3D, causal or not)."""
    if dims == 2:
        return nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=spatial_padding_mode.value,
        )
    elif dims == 3:
        if causal:
            return CausalConv3d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                dilation=dilation,
                groups=groups,
                bias=bias,
                spatial_padding_mode=spatial_padding_mode,
            )
        return nn.Conv3d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=spatial_padding_mode.value,
        )
    else:
        raise ValueError(f"unsupported dimensions: {dims}")

fastvideo.models.vaes.ltx2vae.make_linear_nd

make_linear_nd(dims: int, in_channels: int, out_channels: int, bias: bool = True) -> nn.Module

Create a 1x1 convolution (pointwise linear).

Source
def make_linear_nd(
    dims: int,
    in_channels: int,
    out_channels: int,
    bias: bool = True,
) -> nn.Module:
    """Create a 1x1 convolution (pointwise linear)."""
    if dims == 2:
        return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
    elif dims == 3:
        return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
    else:
        raise ValueError(f"unsupported dimensions: {dims}")

fastvideo.models.vaes.ltx2vae.ResnetBlock3D

ResnetBlock3D(dims: int, in_channels: int, out_channels: int | None = None, dropout: float = 0.0, groups: int = 32, eps: float = 1e-06, norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, inject_noise: bool = False, timestep_conditioning: bool = False, spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS)

A 3D ResNet block with optional timestep conditioning and noise injection.

Source
class ResnetBlock3D(nn.Module):
    """A 3D ResNet block with optional timestep conditioning and noise injection."""

    def __init__(
        self,
        dims: int,
        in_channels: int,
        out_channels: int | None = None,
        dropout: float = 0.0,
        groups: int = 32,
        eps: float = 1e-6,
        norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
        inject_noise: bool = False,
        timestep_conditioning: bool = False,
        spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    ):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.inject_noise = inject_noise

        if norm_layer == NormLayerType.GROUP_NORM:
            self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
        elif norm_layer == NormLayerType.PIXEL_NORM:
            self.norm1 = PixelNorm()

        self.non_linearity = nn.SiLU()

        self.conv1 = make_conv_nd(
            dims,
            in_channels,
            out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            causal=True,
            spatial_padding_mode=spatial_padding_mode,
        )

        if inject_noise:
            self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))

        if norm_layer == NormLayerType.GROUP_NORM:
            self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
        elif norm_layer == NormLayerType.PIXEL_NORM:
            self.norm2 = PixelNorm()

        self.dropout = nn.Dropout(dropout)

        self.conv2 = make_conv_nd(
            dims,
            out_channels,
            out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            causal=True,
            spatial_padding_mode=spatial_padding_mode,
        )

        if inject_noise:
            self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))

        self.conv_shortcut = (make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
                              if in_channels != out_channels else nn.Identity())

        self.norm3 = (nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
                      if in_channels != out_channels else nn.Identity())

        self.timestep_conditioning = timestep_conditioning

        if timestep_conditioning:
            self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)

    def _feed_spatial_noise(
        self,
        hidden_states: torch.Tensor,
        per_channel_scale: torch.Tensor,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        spatial_shape = hidden_states.shape[-2:]
        device = hidden_states.device
        dtype = hidden_states.dtype

        spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None]
        scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
        hidden_states = hidden_states + scaled_noise
        return hidden_states

    def forward(
        self,
        input_tensor: torch.Tensor,
        causal: bool = True,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        hidden_states = input_tensor
        batch_size = hidden_states.shape[0]

        hidden_states = self.norm1(hidden_states)
        if self.timestep_conditioning:
            if timestep is None:
                raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
            ada_values = self.scale_shift_table[None, ..., None, None, None].to(
                device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape(
                    batch_size,
                    4,
                    -1,
                    timestep.shape[-3],
                    timestep.shape[-2],
                    timestep.shape[-1],
                )
            shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
            hidden_states = hidden_states * (1 + scale1) + shift1

        hidden_states = self.non_linearity(hidden_states)
        hidden_states = self.conv1(hidden_states, causal=causal)

        if self.inject_noise:
            hidden_states = self._feed_spatial_noise(
                hidden_states,
                self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
                generator=generator,
            )

        hidden_states = self.norm2(hidden_states)

        if self.timestep_conditioning:
            hidden_states = hidden_states * (1 + scale2) + shift2

        hidden_states = self.non_linearity(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.conv2(hidden_states, causal=causal)

        if self.inject_noise:
            hidden_states = self._feed_spatial_noise(
                hidden_states,
                self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
                generator=generator,
            )

        input_tensor = self.norm3(input_tensor)
        input_tensor = self.conv_shortcut(input_tensor)
        output_tensor = input_tensor + hidden_states
        return output_tensor

fastvideo.models.vaes.ltx2vae.ResnetBlock3D.forward

forward(input_tensor: torch.Tensor, causal: bool = True, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> torch.Tensor
Source
def forward(
    self,
    input_tensor: torch.Tensor,
    causal: bool = True,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    hidden_states = input_tensor
    batch_size = hidden_states.shape[0]

    hidden_states = self.norm1(hidden_states)
    if self.timestep_conditioning:
        if timestep is None:
            raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
        ada_values = self.scale_shift_table[None, ..., None, None, None].to(
            device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape(
                batch_size,
                4,
                -1,
                timestep.shape[-3],
                timestep.shape[-2],
                timestep.shape[-1],
            )
        shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
        hidden_states = hidden_states * (1 + scale1) + shift1

    hidden_states = self.non_linearity(hidden_states)
    hidden_states = self.conv1(hidden_states, causal=causal)

    if self.inject_noise:
        hidden_states = self._feed_spatial_noise(
            hidden_states,
            self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
            generator=generator,
        )

    hidden_states = self.norm2(hidden_states)

    if self.timestep_conditioning:
        hidden_states = hidden_states * (1 + scale2) + shift2

    hidden_states = self.non_linearity(hidden_states)
    hidden_states = self.dropout(hidden_states)
    hidden_states = self.conv2(hidden_states, causal=causal)

    if self.inject_noise:
        hidden_states = self._feed_spatial_noise(
            hidden_states,
            self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
            generator=generator,
        )

    input_tensor = self.norm3(input_tensor)
    input_tensor = self.conv_shortcut(input_tensor)
    output_tensor = input_tensor + hidden_states
    return output_tensor

fastvideo.models.vaes.ltx2vae.UNetMidBlock3D

UNetMidBlock3D(dims: int, in_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_groups: int = 32, norm_layer: NormLayerType = NormLayerType.GROUP_NORM, inject_noise: bool = False, timestep_conditioning: bool = False, spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, attention_head_dim: int | None = None)

A 3D UNet mid-block with multiple residual blocks.

Source
class UNetMidBlock3D(nn.Module):
    """A 3D UNet mid-block with multiple residual blocks."""

    def __init__(
            self,
            dims: int,
            in_channels: int,
            dropout: float = 0.0,
            num_layers: int = 1,
            resnet_eps: float = 1e-6,
            resnet_groups: int = 32,
            norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
            inject_noise: bool = False,
            timestep_conditioning: bool = False,
            spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
            attention_head_dim: int | None = None,  # unused, for compatibility
    ):
        super().__init__()
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

        self.timestep_conditioning = timestep_conditioning

        if timestep_conditioning:
            self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=in_channels * 4,
                                                                           size_emb_dim=0)

        self.res_blocks = nn.ModuleList([
            ResnetBlock3D(
                dims=dims,
                in_channels=in_channels,
                out_channels=in_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=dropout,
                norm_layer=norm_layer,
                inject_noise=inject_noise,
                timestep_conditioning=timestep_conditioning,
                spatial_padding_mode=spatial_padding_mode,
            ) for _ in range(num_layers)
        ])

    def forward(
        self,
        hidden_states: torch.Tensor,
        causal: bool = True,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        timestep_embed = None
        if self.timestep_conditioning:
            if timestep is None:
                raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
            batch_size = hidden_states.shape[0]
            timestep_embed = self.time_embedder(
                timestep=timestep.flatten(),
                hidden_dtype=hidden_states.dtype,
            )
            timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)

        for resnet in self.res_blocks:
            hidden_states = resnet(
                hidden_states,
                causal=causal,
                timestep=timestep_embed,
                generator=generator,
            )

        return hidden_states

fastvideo.models.vaes.ltx2vae.UNetMidBlock3D.forward

forward(hidden_states: torch.Tensor, causal: bool = True, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    causal: bool = True,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    timestep_embed = None
    if self.timestep_conditioning:
        if timestep is None:
            raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
        batch_size = hidden_states.shape[0]
        timestep_embed = self.time_embedder(
            timestep=timestep.flatten(),
            hidden_dtype=hidden_states.dtype,
        )
        timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)

    for resnet in self.res_blocks:
        hidden_states = resnet(
            hidden_states,
            causal=causal,
            timestep=timestep_embed,
            generator=generator,
        )

    return hidden_states

fastvideo.models.vaes.ltx2vae.PixArtAlphaCombinedTimestepSizeEmbeddings

PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim: int, size_emb_dim: int = 0)

Timestep embeddings for decoder conditioning.

Source
class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module):
    """Timestep embeddings for decoder conditioning."""

    def __init__(self, embedding_dim: int, size_emb_dim: int = 0):
        super().__init__()
        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)

    def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor:
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
        return timesteps_emb

fastvideo.models.vaes.ltx2vae.PixArtAlphaCombinedTimestepSizeEmbeddings.forward

forward(timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor
Source
def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor:
    timesteps_proj = self.time_proj(timestep)
    timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
    return timesteps_emb

fastvideo.models.vaes.ltx2vae.Timesteps

Timesteps(num_channels: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0)

Sinusoidal timestep embeddings.

Source
class Timesteps(nn.Module):
    """Sinusoidal timestep embeddings."""

    def __init__(self, num_channels: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0):
        super().__init__()
        self.num_channels = num_channels
        self.flip_sin_to_cos = flip_sin_to_cos
        self.downscale_freq_shift = downscale_freq_shift

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        half_dim = self.num_channels // 2
        exponent = -math.log(10000) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
        exponent = exponent / (half_dim - self.downscale_freq_shift)
        emb = torch.exp(exponent)
        emb = timesteps[:, None].float() * emb[None, :]
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
        if self.flip_sin_to_cos:
            emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
        return emb

fastvideo.models.vaes.ltx2vae.Timesteps.forward

forward(timesteps: torch.Tensor) -> torch.Tensor
Source
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
    half_dim = self.num_channels // 2
    exponent = -math.log(10000) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
    exponent = exponent / (half_dim - self.downscale_freq_shift)
    emb = torch.exp(exponent)
    emb = timesteps[:, None].float() * emb[None, :]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
    if self.flip_sin_to_cos:
        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
    return emb

fastvideo.models.vaes.ltx2vae.TimestepEmbedding

TimestepEmbedding(in_channels: int, time_embed_dim: int)

MLP for timestep embeddings.

Source
class TimestepEmbedding(nn.Module):
    """MLP for timestep embeddings."""

    def __init__(self, in_channels: int, time_embed_dim: int):
        super().__init__()
        self.linear_1 = nn.Linear(in_channels, time_embed_dim)
        self.act = nn.SiLU()
        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)

    def forward(self, sample: torch.Tensor) -> torch.Tensor:
        sample = self.linear_1(sample)
        sample = self.act(sample)
        sample = self.linear_2(sample)
        return sample

fastvideo.models.vaes.ltx2vae.TimestepEmbedding.forward

forward(sample: torch.Tensor) -> torch.Tensor
Source
def forward(self, sample: torch.Tensor) -> torch.Tensor:
    sample = self.linear_1(sample)
    sample = self.act(sample)
    sample = self.linear_2(sample)
    return sample

fastvideo.models.vaes.ltx2vae.SpaceToDepthDownsample

SpaceToDepthDownsample(dims: int, in_channels: int, out_channels: int, stride: Tuple[int, int, int], spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS)

Downsampling via space-to-depth with residual connection.

Source
class SpaceToDepthDownsample(nn.Module):
    """Downsampling via space-to-depth with residual connection."""

    def __init__(
        self,
        dims: int,
        in_channels: int,
        out_channels: int,
        stride: Tuple[int, int, int],
        spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    ):
        super().__init__()
        self.stride = stride
        self.group_size = in_channels * math.prod(stride) // out_channels
        self.conv = make_conv_nd(
            dims=dims,
            in_channels=in_channels,
            out_channels=out_channels // math.prod(stride),
            kernel_size=3,
            stride=1,
            causal=True,
            spatial_padding_mode=spatial_padding_mode,
        )

    def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
        if self.stride[0] == 2:
            x = torch.cat([x[:, :, :1, :, :], x], dim=2)

        x_in = rearrange(
            x,
            "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
            p1=self.stride[0],
            p2=self.stride[1],
            p3=self.stride[2],
        )
        x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
        x_in = x_in.mean(dim=2)

        x = self.conv(x, causal=causal)
        x = rearrange(
            x,
            "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
            p1=self.stride[0],
            p2=self.stride[1],
            p3=self.stride[2],
        )

        x = x + x_in
        return x

fastvideo.models.vaes.ltx2vae.SpaceToDepthDownsample.forward

forward(x: torch.Tensor, causal: bool = True) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
    if self.stride[0] == 2:
        x = torch.cat([x[:, :, :1, :, :], x], dim=2)

    x_in = rearrange(
        x,
        "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
        p1=self.stride[0],
        p2=self.stride[1],
        p3=self.stride[2],
    )
    x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
    x_in = x_in.mean(dim=2)

    x = self.conv(x, causal=causal)
    x = rearrange(
        x,
        "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
        p1=self.stride[0],
        p2=self.stride[1],
        p3=self.stride[2],
    )

    x = x + x_in
    return x

fastvideo.models.vaes.ltx2vae.DepthToSpaceUpsample

DepthToSpaceUpsample(dims: int, in_channels: int, stride: Tuple[int, int, int], residual: bool = False, out_channels_reduction_factor: int = 1, spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS)

Upsampling via depth-to-space (pixel shuffle).

Source
class DepthToSpaceUpsample(nn.Module):
    """Upsampling via depth-to-space (pixel shuffle)."""

    def __init__(
        self,
        dims: int,
        in_channels: int,
        stride: Tuple[int, int, int],
        residual: bool = False,
        out_channels_reduction_factor: int = 1,
        spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    ):
        super().__init__()
        self.stride = stride
        self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
        self.conv = make_conv_nd(
            dims=dims,
            in_channels=in_channels,
            out_channels=self.out_channels,
            kernel_size=3,
            stride=1,
            causal=True,
            spatial_padding_mode=spatial_padding_mode,
        )
        self.residual = residual
        self.out_channels_reduction_factor = out_channels_reduction_factor

    def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
        if self.residual:
            x_in = rearrange(
                x,
                "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
                p1=self.stride[0],
                p2=self.stride[1],
                p3=self.stride[2],
            )
            num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
            x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
            if self.stride[0] == 2:
                x_in = x_in[:, :, 1:, :, :]

        x = self.conv(x, causal=causal)
        x = rearrange(
            x,
            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
            p1=self.stride[0],
            p2=self.stride[1],
            p3=self.stride[2],
        )
        if self.stride[0] == 2:
            x = x[:, :, 1:, :, :]
        if self.residual:
            x = x + x_in
        return x

fastvideo.models.vaes.ltx2vae.DepthToSpaceUpsample.forward

forward(x: torch.Tensor, causal: bool = True) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
    if self.residual:
        x_in = rearrange(
            x,
            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
            p1=self.stride[0],
            p2=self.stride[1],
            p3=self.stride[2],
        )
        num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
        x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
        if self.stride[0] == 2:
            x_in = x_in[:, :, 1:, :, :]

    x = self.conv(x, causal=causal)
    x = rearrange(
        x,
        "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
        p1=self.stride[0],
        p2=self.stride[1],
        p3=self.stride[2],
    )
    if self.stride[0] == 2:
        x = x[:, :, 1:, :, :]
    if self.residual:
        x = x + x_in
    return x

fastvideo.models.vaes.ltx2vae.VideoEncoder

VideoEncoder(convolution_dimensions: int = 3, in_channels: int = 3, out_channels: int = 128, encoder_blocks: list[tuple[str, int]] | list[tuple[str, dict[str, Any]]] = [], patch_size: int = 4, norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, latent_log_var: LogVarianceType = LogVarianceType.UNIFORM, encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS)

LTX-2 Video Encoder. Encodes video frames into a latent representation.

Source
class VideoEncoder(nn.Module):
    """
    LTX-2 Video Encoder. Encodes video frames into a latent representation.
    """

    _DEFAULT_NORM_NUM_GROUPS = 32

    def __init__(
        self,
        convolution_dimensions: int = 3,
        in_channels: int = 3,
        out_channels: int = 128,
        encoder_blocks: list[tuple[str, int]] | list[tuple[str, dict[str, Any]]] = [],
        patch_size: int = 4,
        norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
        latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
        encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    ):
        super().__init__()

        self.patch_size = patch_size
        self.norm_layer = norm_layer
        self.latent_channels = out_channels
        self.latent_log_var = latent_log_var
        self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS

        self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)

        in_channels = in_channels * patch_size**2
        feature_channels = out_channels

        self.conv_in = make_conv_nd(
            dims=convolution_dimensions,
            in_channels=in_channels,
            out_channels=feature_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            causal=True,
            spatial_padding_mode=encoder_spatial_padding_mode,
        )

        self.down_blocks = nn.ModuleList([])

        for block_name, block_params in encoder_blocks:
            block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params

            block, feature_channels = _make_encoder_block(
                block_name=block_name,
                block_config=block_config,
                in_channels=feature_channels,
                convolution_dimensions=convolution_dimensions,
                norm_layer=norm_layer,
                norm_num_groups=self._norm_num_groups,
                spatial_padding_mode=encoder_spatial_padding_mode,
            )

            self.down_blocks.append(block)

        if norm_layer == NormLayerType.GROUP_NORM:
            self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
        elif norm_layer == NormLayerType.PIXEL_NORM:
            self.conv_norm_out = PixelNorm()

        self.conv_act = nn.SiLU()

        conv_out_channels = out_channels
        if latent_log_var == LogVarianceType.PER_CHANNEL:
            conv_out_channels *= 2
        elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
            conv_out_channels += 1
        elif latent_log_var != LogVarianceType.NONE:
            raise ValueError(f"Invalid latent_log_var: {latent_log_var}")

        self.conv_out = make_conv_nd(
            dims=convolution_dimensions,
            in_channels=feature_channels,
            out_channels=conv_out_channels,
            kernel_size=3,
            padding=1,
            causal=True,
            spatial_padding_mode=encoder_spatial_padding_mode,
        )

    def forward(self, sample: torch.Tensor) -> torch.Tensor:
        frames_count = sample.shape[2]
        if ((frames_count - 1) % 8) != 0:
            raise ValueError("Invalid number of frames: Encode input must have 1 + 8 * x frames "
                             "(e.g., 1, 9, 17, ...). Please check your input.")

        sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
        sample = self.conv_in(sample)

        for down_block in self.down_blocks:
            sample = down_block(sample)

        sample = self.conv_norm_out(sample)
        sample = self.conv_act(sample)
        sample = self.conv_out(sample)

        if self.latent_log_var == LogVarianceType.UNIFORM:
            if sample.shape[1] < 2:
                raise ValueError(f"Invalid channel count for UNIFORM mode: expected at least 2, got {sample.shape[1]}")
            means = sample[:, :-1, ...]
            logvar = sample[:, -1:, ...]
            num_channels = means.shape[1]
            repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
            repeated_logvar = logvar.repeat(*repeat_shape)
            sample = torch.cat([means, repeated_logvar], dim=1)
        elif self.latent_log_var == LogVarianceType.CONSTANT:
            sample = sample[:, :-1, ...]
            approx_ln_0 = -30
            sample = torch.cat(
                [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
                dim=1,
            )

        means, _ = torch.chunk(sample, 2, dim=1)
        return self.per_channel_statistics.normalize(means)

fastvideo.models.vaes.ltx2vae.VideoEncoder.forward

forward(sample: torch.Tensor) -> torch.Tensor
Source
def forward(self, sample: torch.Tensor) -> torch.Tensor:
    frames_count = sample.shape[2]
    if ((frames_count - 1) % 8) != 0:
        raise ValueError("Invalid number of frames: Encode input must have 1 + 8 * x frames "
                         "(e.g., 1, 9, 17, ...). Please check your input.")

    sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
    sample = self.conv_in(sample)

    for down_block in self.down_blocks:
        sample = down_block(sample)

    sample = self.conv_norm_out(sample)
    sample = self.conv_act(sample)
    sample = self.conv_out(sample)

    if self.latent_log_var == LogVarianceType.UNIFORM:
        if sample.shape[1] < 2:
            raise ValueError(f"Invalid channel count for UNIFORM mode: expected at least 2, got {sample.shape[1]}")
        means = sample[:, :-1, ...]
        logvar = sample[:, -1:, ...]
        num_channels = means.shape[1]
        repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
        repeated_logvar = logvar.repeat(*repeat_shape)
        sample = torch.cat([means, repeated_logvar], dim=1)
    elif self.latent_log_var == LogVarianceType.CONSTANT:
        sample = sample[:, :-1, ...]
        approx_ln_0 = -30
        sample = torch.cat(
            [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
            dim=1,
        )

    means, _ = torch.chunk(sample, 2, dim=1)
    return self.per_channel_statistics.normalize(means)

fastvideo.models.vaes.ltx2vae.VideoDecoder

VideoDecoder(convolution_dimensions: int = 3, in_channels: int = 128, out_channels: int = 3, decoder_blocks: list[tuple[str, int | dict]] = [], patch_size: int = 4, norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, causal: bool = False, timestep_conditioning: bool = False, decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, base_channels: int = 128)

LTX-2 Video Decoder. Decodes latent representation into video frames.

Source
class VideoDecoder(nn.Module):
    """
    LTX-2 Video Decoder. Decodes latent representation into video frames.
    """

    _DEFAULT_NORM_NUM_GROUPS = 32

    def __init__(
        self,
        convolution_dimensions: int = 3,
        in_channels: int = 128,
        out_channels: int = 3,
        decoder_blocks: list[tuple[str, int | dict]] = [],
        patch_size: int = 4,
        norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
        causal: bool = False,
        timestep_conditioning: bool = False,
        decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
        base_channels: int = 128,
    ):
        super().__init__()

        self.patch_size = patch_size
        out_channels = out_channels * patch_size**2
        self.causal = causal
        self.timestep_conditioning = timestep_conditioning
        self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS

        self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)

        self.decode_noise_scale = 0.025
        self.decode_timestep = 0.05

        feature_channels = base_channels
        for block_name, block_params in list(reversed(decoder_blocks)):
            block_config = block_params if isinstance(block_params, dict) else {}
            if block_name == "res_x_y":
                feature_channels = feature_channels * block_config.get("multiplier", 2)
            if block_name == "compress_all":
                feature_channels = feature_channels * block_config.get("multiplier", 1)
            if block_name == "compress_space":
                feature_channels = feature_channels * block_config.get("multiplier", 1)
            if block_name == "compress_time":
                feature_channels = feature_channels * block_config.get("multiplier", 1)

        self.conv_in = make_conv_nd(
            dims=convolution_dimensions,
            in_channels=in_channels,
            out_channels=feature_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            causal=True,
            spatial_padding_mode=decoder_spatial_padding_mode,
        )

        self.up_blocks = nn.ModuleList([])

        for block_name, block_params in list(reversed(decoder_blocks)):
            block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params

            block, feature_channels = _make_decoder_block(
                block_name=block_name,
                block_config=block_config,
                in_channels=feature_channels,
                convolution_dimensions=convolution_dimensions,
                norm_layer=norm_layer,
                timestep_conditioning=timestep_conditioning,
                norm_num_groups=self._norm_num_groups,
                spatial_padding_mode=decoder_spatial_padding_mode,
            )

            self.up_blocks.append(block)

        if norm_layer == NormLayerType.GROUP_NORM:
            self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
        elif norm_layer == NormLayerType.PIXEL_NORM:
            self.conv_norm_out = PixelNorm()

        self.conv_act = nn.SiLU()
        self.conv_out = make_conv_nd(
            dims=convolution_dimensions,
            in_channels=feature_channels,
            out_channels=out_channels,
            kernel_size=3,
            padding=1,
            causal=True,
            spatial_padding_mode=decoder_spatial_padding_mode,
        )

        if timestep_conditioning:
            self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
            self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=feature_channels * 2,
                                                                                size_emb_dim=0)
            self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))

    def forward(
        self,
        sample: torch.Tensor,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        batch_size = sample.shape[0]

        if self.timestep_conditioning:
            noise = (torch.randn(
                sample.size(),
                generator=generator,
                dtype=sample.dtype,
                device=sample.device,
            ) * self.decode_noise_scale)
            sample = noise + (1.0 - self.decode_noise_scale) * sample

        sample = self.per_channel_statistics.un_normalize(sample)

        if timestep is None and self.timestep_conditioning:
            timestep = torch.full((batch_size, ), self.decode_timestep, device=sample.device, dtype=sample.dtype)

        sample = self.conv_in(sample, causal=self.causal)

        scaled_timestep = None
        if self.timestep_conditioning:
            if timestep is None:
                raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
            scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)

        for up_block in self.up_blocks:
            if isinstance(up_block, UNetMidBlock3D):
                block_kwargs = {
                    "causal": self.causal,
                    "timestep": scaled_timestep if self.timestep_conditioning else None,
                    "generator": generator,
                }
                sample = up_block(sample, **block_kwargs)
            elif isinstance(up_block, ResnetBlock3D):
                sample = up_block(sample, causal=self.causal, generator=generator)
            else:
                sample = up_block(sample, causal=self.causal)

        sample = self.conv_norm_out(sample)

        if self.timestep_conditioning:
            embedded_timestep = self.last_time_embedder(
                timestep=scaled_timestep.flatten(),
                hidden_dtype=sample.dtype,
            )
            embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
            ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
                device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape(
                    batch_size,
                    2,
                    -1,
                    embedded_timestep.shape[-3],
                    embedded_timestep.shape[-2],
                    embedded_timestep.shape[-1],
                )
            shift, scale = ada_values.unbind(dim=1)
            sample = sample * (1 + scale) + shift

        sample = self.conv_act(sample)
        sample = self.conv_out(sample, causal=self.causal)
        sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)

        return sample

fastvideo.models.vaes.ltx2vae.VideoDecoder.forward

forward(sample: torch.Tensor, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> torch.Tensor
Source
def forward(
    self,
    sample: torch.Tensor,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    batch_size = sample.shape[0]

    if self.timestep_conditioning:
        noise = (torch.randn(
            sample.size(),
            generator=generator,
            dtype=sample.dtype,
            device=sample.device,
        ) * self.decode_noise_scale)
        sample = noise + (1.0 - self.decode_noise_scale) * sample

    sample = self.per_channel_statistics.un_normalize(sample)

    if timestep is None and self.timestep_conditioning:
        timestep = torch.full((batch_size, ), self.decode_timestep, device=sample.device, dtype=sample.dtype)

    sample = self.conv_in(sample, causal=self.causal)

    scaled_timestep = None
    if self.timestep_conditioning:
        if timestep is None:
            raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True")
        scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)

    for up_block in self.up_blocks:
        if isinstance(up_block, UNetMidBlock3D):
            block_kwargs = {
                "causal": self.causal,
                "timestep": scaled_timestep if self.timestep_conditioning else None,
                "generator": generator,
            }
            sample = up_block(sample, **block_kwargs)
        elif isinstance(up_block, ResnetBlock3D):
            sample = up_block(sample, causal=self.causal, generator=generator)
        else:
            sample = up_block(sample, causal=self.causal)

    sample = self.conv_norm_out(sample)

    if self.timestep_conditioning:
        embedded_timestep = self.last_time_embedder(
            timestep=scaled_timestep.flatten(),
            hidden_dtype=sample.dtype,
        )
        embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
        ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
            device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape(
                batch_size,
                2,
                -1,
                embedded_timestep.shape[-3],
                embedded_timestep.shape[-2],
                embedded_timestep.shape[-1],
            )
        shift, scale = ada_values.unbind(dim=1)
        sample = sample * (1 + scale) + shift

    sample = self.conv_act(sample)
    sample = self.conv_out(sample, causal=self.causal)
    sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)

    return sample

fastvideo.models.vaes.ltx2vae.VideoEncoderConfigurator

VideoEncoderConfigurator()

Configurator for creating a video VAE Encoder from a configuration dictionary.

Source
class VideoEncoderConfigurator:
    """Configurator for creating a video VAE Encoder from a configuration dictionary."""

    @classmethod
    def from_config(cls, config: dict) -> VideoEncoder:
        config = config.get("vae", config)
        convolution_dimensions = config.get("dims", 3)
        in_channels = config.get("in_channels", 3)
        latent_channels = config.get("latent_channels", 128)
        encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
        encoder_blocks = config.get("encoder_blocks", [])
        patch_size = config.get("patch_size", 4)
        norm_layer_str = config.get("norm_layer", "pixel_norm")
        latent_log_var_str = config.get("latent_log_var", "uniform")

        return VideoEncoder(
            convolution_dimensions=convolution_dimensions,
            in_channels=in_channels,
            out_channels=latent_channels,
            encoder_blocks=encoder_blocks,
            patch_size=patch_size,
            norm_layer=NormLayerType(norm_layer_str),
            latent_log_var=LogVarianceType(latent_log_var_str),
            encoder_spatial_padding_mode=encoder_spatial_padding_mode,
        )

fastvideo.models.vaes.ltx2vae.VideoEncoderConfigurator.from_config

from_config(config: dict) -> VideoEncoder
Source
@classmethod
def from_config(cls, config: dict) -> VideoEncoder:
    config = config.get("vae", config)
    convolution_dimensions = config.get("dims", 3)
    in_channels = config.get("in_channels", 3)
    latent_channels = config.get("latent_channels", 128)
    encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
    encoder_blocks = config.get("encoder_blocks", [])
    patch_size = config.get("patch_size", 4)
    norm_layer_str = config.get("norm_layer", "pixel_norm")
    latent_log_var_str = config.get("latent_log_var", "uniform")

    return VideoEncoder(
        convolution_dimensions=convolution_dimensions,
        in_channels=in_channels,
        out_channels=latent_channels,
        encoder_blocks=encoder_blocks,
        patch_size=patch_size,
        norm_layer=NormLayerType(norm_layer_str),
        latent_log_var=LogVarianceType(latent_log_var_str),
        encoder_spatial_padding_mode=encoder_spatial_padding_mode,
    )

fastvideo.models.vaes.ltx2vae.VideoDecoderConfigurator

VideoDecoderConfigurator()

Configurator for creating a video VAE Decoder from a configuration dictionary.

Source
class VideoDecoderConfigurator:
    """Configurator for creating a video VAE Decoder from a configuration dictionary."""

    @classmethod
    def from_config(cls, config: dict) -> VideoDecoder:
        config = config.get("vae", config)
        convolution_dimensions = config.get("dims", 3)
        latent_channels = config.get("latent_channels", 128)
        decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
        out_channels = config.get("out_channels", 3)
        decoder_blocks = config.get("decoder_blocks", [])
        patch_size = config.get("patch_size", 4)
        norm_layer_str = config.get("norm_layer", "pixel_norm")
        causal = config.get("causal_decoder", False)
        timestep_conditioning = config.get("timestep_conditioning", True)
        base_channels = config.get("decoder_base_channels", 128)

        return VideoDecoder(
            convolution_dimensions=convolution_dimensions,
            in_channels=latent_channels,
            out_channels=out_channels,
            decoder_blocks=decoder_blocks,
            patch_size=patch_size,
            norm_layer=NormLayerType(norm_layer_str),
            causal=causal,
            timestep_conditioning=timestep_conditioning,
            decoder_spatial_padding_mode=decoder_spatial_padding_mode,
            base_channels=base_channels,
        )

fastvideo.models.vaes.ltx2vae.VideoDecoderConfigurator.from_config

from_config(config: dict) -> VideoDecoder
Source
@classmethod
def from_config(cls, config: dict) -> VideoDecoder:
    config = config.get("vae", config)
    convolution_dimensions = config.get("dims", 3)
    latent_channels = config.get("latent_channels", 128)
    decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
    out_channels = config.get("out_channels", 3)
    decoder_blocks = config.get("decoder_blocks", [])
    patch_size = config.get("patch_size", 4)
    norm_layer_str = config.get("norm_layer", "pixel_norm")
    causal = config.get("causal_decoder", False)
    timestep_conditioning = config.get("timestep_conditioning", True)
    base_channels = config.get("decoder_base_channels", 128)

    return VideoDecoder(
        convolution_dimensions=convolution_dimensions,
        in_channels=latent_channels,
        out_channels=out_channels,
        decoder_blocks=decoder_blocks,
        patch_size=patch_size,
        norm_layer=NormLayerType(norm_layer_str),
        causal=causal,
        timestep_conditioning=timestep_conditioning,
        decoder_spatial_padding_mode=decoder_spatial_padding_mode,
        base_channels=base_channels,
    )

fastvideo.models.vaes.ltx2vae.LTX2VideoEncoder

LTX2VideoEncoder(config: dict[str, Any])

LTX-2 Video Encoder wrapper for FastVideo compatibility.

Source
class LTX2VideoEncoder(nn.Module):
    """LTX-2 Video Encoder wrapper for FastVideo compatibility."""

    def __init__(self, config: dict[str, Any]):
        super().__init__()
        self.model: VideoEncoder = VideoEncoderConfigurator.from_config(config)

    def forward(self, sample: torch.Tensor) -> torch.Tensor:
        return self.model(sample)

fastvideo.models.vaes.ltx2vae.LTX2VideoEncoder.forward

forward(sample: torch.Tensor) -> torch.Tensor
Source
def forward(self, sample: torch.Tensor) -> torch.Tensor:
    return self.model(sample)

fastvideo.models.vaes.ltx2vae.LTX2VideoDecoder

LTX2VideoDecoder(config: dict[str, Any])

LTX-2 Video Decoder wrapper for FastVideo compatibility.

Source
class LTX2VideoDecoder(nn.Module):
    """LTX-2 Video Decoder wrapper for FastVideo compatibility."""

    def __init__(self, config: dict[str, Any]):
        super().__init__()
        self.model: VideoDecoder = VideoDecoderConfigurator.from_config(config)

    def forward(
        self,
        sample: torch.Tensor,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        return self.model(sample, timestep=timestep, generator=generator)

fastvideo.models.vaes.ltx2vae.LTX2VideoDecoder.forward

forward(sample: torch.Tensor, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> torch.Tensor
Source
def forward(
    self,
    sample: torch.Tensor,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    return self.model(sample, timestep=timestep, generator=generator)

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder

LTX2CausalVideoAutoencoder(config: dict[str, Any])

LTX-2 VAE that exposes FastVideo's VAE encode/decode interface. Supports tiled decoding to reduce memory usage for high-resolution videos.

Source
class LTX2CausalVideoAutoencoder(nn.Module):
    """
    LTX-2 VAE that exposes FastVideo's VAE encode/decode interface.
    Supports tiled decoding to reduce memory usage for high-resolution videos.
    """

    # LTX-2 VAE scale factors
    TIME_SCALE: int = 8
    SPATIAL_SCALE: int = 32
    _compile_conditions = [_is_ltx2_vae_codec]

    def __init__(self, config: dict[str, Any]):
        super().__init__()
        self.config = config
        self.encoder = VideoEncoderConfigurator.from_config(config)
        self.decoder = VideoDecoderConfigurator.from_config(config)
        self._use_tiling: bool = False
        self._use_channels_last_3d: bool = False

        if envs.FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D.get():
            self.enable_channels_last_3d()

    def _as_channels_last_3d(self, tensor: torch.Tensor) -> torch.Tensor:
        if not self._use_channels_last_3d or tensor.ndim != 5 or tensor.is_contiguous(
                memory_format=torch.channels_last_3d):
            return tensor
        return tensor.contiguous(memory_format=torch.channels_last_3d)

    def enable_channels_last_3d(self) -> None:
        """Enable channels-last layout for 3D VAE convolutions."""
        self._use_channels_last_3d = True
        self.encoder.to(memory_format=torch.channels_last_3d)
        self.decoder.to(memory_format=torch.channels_last_3d)

    def disable_channels_last_3d(self) -> None:
        """Restore contiguous layout for 3D VAE convolutions."""
        self._use_channels_last_3d = False
        self.encoder.to(memory_format=torch.contiguous_format)
        self.decoder.to(memory_format=torch.contiguous_format)

    def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
        x = self._as_channels_last_3d(x)
        means = self.encoder(x)
        zeros = torch.zeros_like(means)
        return DiagonalGaussianDistribution(torch.cat([means, zeros], dim=1), deterministic=True)

    def decode(
        self,
        z: torch.Tensor,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        """Decode latents to video, using tiling if enabled."""
        z = self._as_channels_last_3d(z)
        if self._use_tiling:
            # Collect all chunks from tiled decode and concatenate
            chunks = list(self.tiled_decode(z, TilingConfig.default(), timestep, generator))
            return torch.cat(chunks, dim=2)  # Concatenate along temporal dimension
        return self.decoder(z, timestep=timestep, generator=generator)

    def enable_tiling(self) -> None:
        """Enable tiled decoding with default configuration."""
        self._use_tiling = True

    def disable_tiling(self) -> None:
        """Disable tiled decoding."""
        self._use_tiling = False

    def _prepare_tiles(
        self,
        latent: torch.Tensor,
        tiling_config: TilingConfig | None,
    ) -> List[Tile]:
        """Prepare tiles for tiled decoding based on tiling configuration."""
        splitters: List[SplitOperation] = [DEFAULT_SPLIT_OPERATION] * 5
        mappers: List[MappingOperation] = [DEFAULT_MAPPING_OPERATION] * 5

        if tiling_config is not None and tiling_config.spatial_config is not None:
            cfg = tiling_config.spatial_config
            tile_size = cfg.tile_size_in_pixels // self.SPATIAL_SCALE
            overlap = cfg.tile_overlap_in_pixels // self.SPATIAL_SCALE
            splitters[3] = split_in_spatial(tile_size, overlap)
            splitters[4] = split_in_spatial(tile_size, overlap)
            mappers[3] = to_mapping_operation(map_spatial_slice, self.SPATIAL_SCALE)
            mappers[4] = to_mapping_operation(map_spatial_slice, self.SPATIAL_SCALE)

        if tiling_config is not None and tiling_config.temporal_config is not None:
            cfg = tiling_config.temporal_config
            tile_size = cfg.tile_size_in_frames // self.TIME_SCALE
            overlap = cfg.tile_overlap_in_frames // self.TIME_SCALE
            splitters[2] = split_in_temporal(tile_size, overlap)
            mappers[2] = to_mapping_operation(map_temporal_slice, self.TIME_SCALE)

        return create_tiles(latent.shape, splitters, mappers)

    def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
        """Group tiles by their temporal output slice."""
        if not tiles:
            return []

        groups = []
        current_slice = tiles[0].out_coords[2]
        current_group = []

        for tile in tiles:
            tile_slice = tile.out_coords[2]
            if tile_slice == current_slice:
                current_group.append(tile)
            else:
                groups.append(current_group)
                current_slice = tile_slice
                current_group = [tile]

        if current_group:
            groups.append(current_group)

        return groups

    def _accumulate_temporal_group_into_buffer(
        self,
        group_tiles: List[Tile],
        buffer: torch.Tensor,
        latent: torch.Tensor,
        timestep: torch.Tensor | None,
        generator: torch.Generator | None,
    ) -> torch.Tensor:
        """Decode and accumulate all tiles of a temporal group into a local buffer."""
        temporal_slice = group_tiles[0].out_coords[2]
        weights = torch.zeros_like(buffer)

        for tile in group_tiles:
            decoded_tile = self.decoder(latent[tile.in_coords], timestep, generator)
            mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype)
            temporal_offset = tile.out_coords[2].start - temporal_slice.start
            expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start
            decoded_temporal_len = decoded_tile.shape[2]

            actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset)

            chunk_coords = (
                slice(None),  # batch
                slice(None),  # channels
                slice(temporal_offset, temporal_offset + actual_temporal_len),
                tile.out_coords[3],  # height
                tile.out_coords[4],  # width
            )

            decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :]
            mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask

            buffer[chunk_coords] += decoded_slice * mask_slice
            weights[chunk_coords] += mask_slice

        return weights

    def tiled_decode(
        self,
        latent: torch.Tensor,
        tiling_config: TilingConfig | None = None,
        timestep: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
    ) -> Iterator[torch.Tensor]:
        """
        Decode a latent tensor into video frames using tiled processing.
        Splits the latent tensor into tiles, decodes each tile individually,
        and yields video chunks as they become available.

        Args:
            latent: Input latent tensor (B, C, F', H', W').
            tiling_config: Tiling configuration for the latent tensor.
            timestep: Optional timestep for decoder conditioning.
            generator: Optional random generator for deterministic decoding.

        Yields:
            Video chunks (B, C, T, H, W) by temporal slices.
        """
        full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.TIME_SCALE, self.SPATIAL_SCALE)
        tiles = self._prepare_tiles(latent, tiling_config)
        temporal_groups = self._group_tiles_by_temporal_slice(tiles)

        previous_chunk = None
        previous_weights = None
        previous_temporal_slice = None

        for temporal_group_tiles in temporal_groups:
            curr_temporal_slice = temporal_group_tiles[0].out_coords[2]

            temporal_tile_buffer_shape = full_video_shape._replace(frames=curr_temporal_slice.stop -
                                                                   curr_temporal_slice.start, )

            buffer = torch.zeros(
                temporal_tile_buffer_shape.to_torch_shape(),
                device=latent.device,
                dtype=latent.dtype,
            )

            curr_weights = self._accumulate_temporal_group_into_buffer(
                group_tiles=temporal_group_tiles,
                buffer=buffer,
                latent=latent,
                timestep=timestep,
                generator=generator,
            )

            # Blend with previous temporal chunk if it exists
            if previous_chunk is not None:
                if previous_temporal_slice.stop > curr_temporal_slice.start:
                    overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
                    temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)

                    previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
                    previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[:, :,
                                                                                         slice(0, overlap_len), :, :]

                    buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
                    curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[:, :,
                                                                                       temporal_overlap_slice, :, :]

                # Yield the non-overlapping part of the previous chunk
                previous_weights = previous_weights.clamp(min=1e-8)
                yield_len = curr_temporal_slice.start - previous_temporal_slice.start
                yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]

            # Update state for next iteration
            previous_chunk = buffer
            previous_weights = curr_weights
            previous_temporal_slice = curr_temporal_slice

        # Yield any remaining chunk
        if previous_chunk is not None:
            previous_weights = previous_weights.clamp(min=1e-8)
            yield previous_chunk / previous_weights

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.enable_channels_last_3d

enable_channels_last_3d() -> None

Enable channels-last layout for 3D VAE convolutions.

Source
def enable_channels_last_3d(self) -> None:
    """Enable channels-last layout for 3D VAE convolutions."""
    self._use_channels_last_3d = True
    self.encoder.to(memory_format=torch.channels_last_3d)
    self.decoder.to(memory_format=torch.channels_last_3d)

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.disable_channels_last_3d

disable_channels_last_3d() -> None

Restore contiguous layout for 3D VAE convolutions.

Source
def disable_channels_last_3d(self) -> None:
    """Restore contiguous layout for 3D VAE convolutions."""
    self._use_channels_last_3d = False
    self.encoder.to(memory_format=torch.contiguous_format)
    self.decoder.to(memory_format=torch.contiguous_format)

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.encode

encode(x: torch.Tensor) -> DiagonalGaussianDistribution
Source
def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
    x = self._as_channels_last_3d(x)
    means = self.encoder(x)
    zeros = torch.zeros_like(means)
    return DiagonalGaussianDistribution(torch.cat([means, zeros], dim=1), deterministic=True)

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.decode

decode(z: torch.Tensor, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> torch.Tensor

Decode latents to video, using tiling if enabled.

Source
def decode(
    self,
    z: torch.Tensor,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Decode latents to video, using tiling if enabled."""
    z = self._as_channels_last_3d(z)
    if self._use_tiling:
        # Collect all chunks from tiled decode and concatenate
        chunks = list(self.tiled_decode(z, TilingConfig.default(), timestep, generator))
        return torch.cat(chunks, dim=2)  # Concatenate along temporal dimension
    return self.decoder(z, timestep=timestep, generator=generator)

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.enable_tiling

enable_tiling() -> None

Enable tiled decoding with default configuration.

Source
def enable_tiling(self) -> None:
    """Enable tiled decoding with default configuration."""
    self._use_tiling = True

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.disable_tiling

disable_tiling() -> None

Disable tiled decoding.

Source
def disable_tiling(self) -> None:
    """Disable tiled decoding."""
    self._use_tiling = False

fastvideo.models.vaes.ltx2vae.LTX2CausalVideoAutoencoder.tiled_decode

tiled_decode(latent: torch.Tensor, tiling_config: TilingConfig | None = None, timestep: torch.Tensor | None = None, generator: torch.Generator | None = None) -> Iterator[torch.Tensor]

Decode a latent tensor into video frames using tiled processing. Splits the latent tensor into tiles, decodes each tile individually, and yields video chunks as they become available.

Parameters

latent: torch.Tensor

Input latent tensor (B, C, F', H', W').

tiling_config: TilingConfig | None

Tiling configuration for the latent tensor.

timestep: torch.Tensor | None

Optional timestep for decoder conditioning.

generator: torch.Generator | None

Optional random generator for deterministic decoding.

Yields

torch.Tensor

Video chunks (B, C, T, H, W) by temporal slices.

Source
def tiled_decode(
    self,
    latent: torch.Tensor,
    tiling_config: TilingConfig | None = None,
    timestep: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> Iterator[torch.Tensor]:
    """
    Decode a latent tensor into video frames using tiled processing.
    Splits the latent tensor into tiles, decodes each tile individually,
    and yields video chunks as they become available.

    Args:
        latent: Input latent tensor (B, C, F', H', W').
        tiling_config: Tiling configuration for the latent tensor.
        timestep: Optional timestep for decoder conditioning.
        generator: Optional random generator for deterministic decoding.

    Yields:
        Video chunks (B, C, T, H, W) by temporal slices.
    """
    full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.TIME_SCALE, self.SPATIAL_SCALE)
    tiles = self._prepare_tiles(latent, tiling_config)
    temporal_groups = self._group_tiles_by_temporal_slice(tiles)

    previous_chunk = None
    previous_weights = None
    previous_temporal_slice = None

    for temporal_group_tiles in temporal_groups:
        curr_temporal_slice = temporal_group_tiles[0].out_coords[2]

        temporal_tile_buffer_shape = full_video_shape._replace(frames=curr_temporal_slice.stop -
                                                               curr_temporal_slice.start, )

        buffer = torch.zeros(
            temporal_tile_buffer_shape.to_torch_shape(),
            device=latent.device,
            dtype=latent.dtype,
        )

        curr_weights = self._accumulate_temporal_group_into_buffer(
            group_tiles=temporal_group_tiles,
            buffer=buffer,
            latent=latent,
            timestep=timestep,
            generator=generator,
        )

        # Blend with previous temporal chunk if it exists
        if previous_chunk is not None:
            if previous_temporal_slice.stop > curr_temporal_slice.start:
                overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
                temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)

                previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
                previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[:, :,
                                                                                     slice(0, overlap_len), :, :]

                buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
                curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[:, :,
                                                                                   temporal_overlap_slice, :, :]

            # Yield the non-overlapping part of the previous chunk
            previous_weights = previous_weights.clamp(min=1e-8)
            yield_len = curr_temporal_slice.start - previous_temporal_slice.start
            yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]

        # Update state for next iteration
        previous_chunk = buffer
        previous_weights = curr_weights
        previous_temporal_slice = curr_temporal_slice

    # Yield any remaining chunk
    if previous_chunk is not None:
        previous_weights = previous_weights.clamp(min=1e-8)
        yield previous_chunk / previous_weights