v5 release with triple architecture support and prompt enhancer

ltx_video/models/autoencoders/causal_conv3d.py

@@ -0,0 +1,63 @@
+from typing import Tuple, Union
+
+import torch
+import torch.nn as nn
+
+
+class CausalConv3d(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: int = 3,
+        stride: Union[int, Tuple[int]] = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        spatial_padding_mode: str = "zeros",
+        **kwargs,
+    ):
+        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,
+            groups=groups,
+        )
+
+    def forward(self, x, causal: bool = True):
+        if causal:
+            first_frame_pad = x[:, :, :1, :, :].repeat(
+                (1, 1, self.time_kernel_size - 1, 1, 1)
+            )
+            x = torch.concatenate((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.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
+        x = self.conv(x)
+        return x
+
+    @property
+    def weight(self):
+        return self.conv.weight

ltx_video/models/autoencoders/conv_nd_factory.py

@@ -0,0 +1,90 @@
+from typing import Tuple, Union
+
+import torch
+
+from ltx_video.models.autoencoders.dual_conv3d import DualConv3d
+from ltx_video.models.autoencoders.causal_conv3d import CausalConv3d
+
+
+def make_conv_nd(
+    dims: Union[int, Tuple[int, int]],
+    in_channels: int,
+    out_channels: int,
+    kernel_size: int,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    bias=True,
+    causal=False,
+    spatial_padding_mode="zeros",
+    temporal_padding_mode="zeros",
+):
+    if not (spatial_padding_mode == temporal_padding_mode or causal):
+        raise NotImplementedError("spatial and temporal padding modes must be equal")
+    if dims == 2:
+        return torch.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,
+        )
+    elif dims == 3:
+        if causal:
+            return CausalConv3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+                spatial_padding_mode=spatial_padding_mode,
+            )
+        return torch.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,
+        )
+    elif dims == (2, 1):
+        return DualConv3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+            padding_mode=spatial_padding_mode,
+        )
+    else:
+        raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def make_linear_nd(
+    dims: int,
+    in_channels: int,
+    out_channels: int,
+    bias=True,
+):
+    if dims == 2:
+        return torch.nn.Conv2d(
+            in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
+        )
+    elif dims == 3 or dims == (2, 1):
+        return torch.nn.Conv3d(
+            in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
+        )
+    else:
+        raise ValueError(f"unsupported dimensions: {dims}")

ltx_video/models/autoencoders/dual_conv3d.py

@@ -0,0 +1,217 @@
+import math
+from typing import Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class DualConv3d(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        padding: Union[int, Tuple[int, int, int]] = 0,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
+        super(DualConv3d, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.padding_mode = padding_mode
+        # Ensure kernel_size, stride, padding, and dilation are tuples of length 3
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size, kernel_size, kernel_size)
+        if kernel_size == (1, 1, 1):
+            raise ValueError(
+                "kernel_size must be greater than 1. Use make_linear_nd instead."
+            )
+        if isinstance(stride, int):
+            stride = (stride, stride, stride)
+        if isinstance(padding, int):
+            padding = (padding, padding, padding)
+        if isinstance(dilation, int):
+            dilation = (dilation, dilation, dilation)
+
+        # Set parameters for convolutions
+        self.groups = groups
+        self.bias = bias
+
+        # Define the size of the channels after the first convolution
+        intermediate_channels = (
+            out_channels if in_channels < out_channels else in_channels
+        )
+
+        # Define parameters for the first convolution
+        self.weight1 = nn.Parameter(
+            torch.Tensor(
+                intermediate_channels,
+                in_channels // groups,
+                1,
+                kernel_size[1],
+                kernel_size[2],
+            )
+        )
+        self.stride1 = (1, stride[1], stride[2])
+        self.padding1 = (0, padding[1], padding[2])
+        self.dilation1 = (1, dilation[1], dilation[2])
+        if bias:
+            self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
+        else:
+            self.register_parameter("bias1", None)
+
+        # Define parameters for the second convolution
+        self.weight2 = nn.Parameter(
+            torch.Tensor(
+                out_channels, intermediate_channels // groups, kernel_size[0], 1, 1
+            )
+        )
+        self.stride2 = (stride[0], 1, 1)
+        self.padding2 = (padding[0], 0, 0)
+        self.dilation2 = (dilation[0], 1, 1)
+        if bias:
+            self.bias2 = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter("bias2", None)
+
+        # Initialize weights and biases
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight1, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.weight2, a=math.sqrt(5))
+        if self.bias:
+            fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
+            bound1 = 1 / math.sqrt(fan_in1)
+            nn.init.uniform_(self.bias1, -bound1, bound1)
+            fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
+            bound2 = 1 / math.sqrt(fan_in2)
+            nn.init.uniform_(self.bias2, -bound2, bound2)
+
+    def forward(self, x, use_conv3d=False, skip_time_conv=False):
+        if use_conv3d:
+            return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
+        else:
+            return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
+
+    def forward_with_3d(self, x, skip_time_conv):
+        # First convolution
+        x = F.conv3d(
+            x,
+            self.weight1,
+            self.bias1,
+            self.stride1,
+            self.padding1,
+            self.dilation1,
+            self.groups,
+            padding_mode=self.padding_mode,
+        )
+
+        if skip_time_conv:
+            return x
+
+        # Second convolution
+        x = F.conv3d(
+            x,
+            self.weight2,
+            self.bias2,
+            self.stride2,
+            self.padding2,
+            self.dilation2,
+            self.groups,
+            padding_mode=self.padding_mode,
+        )
+
+        return x
+
+    def forward_with_2d(self, x, skip_time_conv):
+        b, c, d, h, w = x.shape
+
+        # First 2D convolution
+        x = rearrange(x, "b c d h w -> (b d) c h w")
+        # Squeeze the depth dimension out of weight1 since it's 1
+        weight1 = self.weight1.squeeze(2)
+        # Select stride, padding, and dilation for the 2D convolution
+        stride1 = (self.stride1[1], self.stride1[2])
+        padding1 = (self.padding1[1], self.padding1[2])
+        dilation1 = (self.dilation1[1], self.dilation1[2])
+        x = F.conv2d(
+            x,
+            weight1,
+            self.bias1,
+            stride1,
+            padding1,
+            dilation1,
+            self.groups,
+            padding_mode=self.padding_mode,
+        )
+
+        _, _, h, w = x.shape
+
+        if skip_time_conv:
+            x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
+            return x
+
+        # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
+        x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
+
+        # Reshape weight2 to match the expected dimensions for conv1d
+        weight2 = self.weight2.squeeze(-1).squeeze(-1)
+        # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
+        stride2 = self.stride2[0]
+        padding2 = self.padding2[0]
+        dilation2 = self.dilation2[0]
+        x = F.conv1d(
+            x,
+            weight2,
+            self.bias2,
+            stride2,
+            padding2,
+            dilation2,
+            self.groups,
+            padding_mode=self.padding_mode,
+        )
+        x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
+
+        return x
+
+    @property
+    def weight(self):
+        return self.weight2
+
+
+def test_dual_conv3d_consistency():
+    # Initialize parameters
+    in_channels = 3
+    out_channels = 5
+    kernel_size = (3, 3, 3)
+    stride = (2, 2, 2)
+    padding = (1, 1, 1)
+
+    # Create an instance of the DualConv3d class
+    dual_conv3d = DualConv3d(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=padding,
+        bias=True,
+    )
+
+    # Example input tensor
+    test_input = torch.randn(1, 3, 10, 10, 10)
+
+    # Perform forward passes with both 3D and 2D settings
+    output_conv3d = dual_conv3d(test_input, use_conv3d=True)
+    output_2d = dual_conv3d(test_input, use_conv3d=False)
+
+    # Assert that the outputs from both methods are sufficiently close
+    assert torch.allclose(
+        output_conv3d, output_2d, atol=1e-6
+    ), "Outputs are not consistent between 3D and 2D convolutions."

ltx_video/models/autoencoders/latent_upsampler.py

@@ -0,0 +1,203 @@
+from typing import Optional, Union
+from pathlib import Path
+import os
+import json
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from diffusers import ConfigMixin, ModelMixin
+from safetensors.torch import safe_open
+
+from ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
+    ):
+        super().__init__()
+        if mid_channels is None:
+            mid_channels = channels
+
+        Conv = nn.Conv2d if dims == 2 else nn.Conv3d
+
+        self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)
+        self.norm1 = nn.GroupNorm(32, mid_channels)
+        self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)
+        self.norm2 = nn.GroupNorm(32, channels)
+        self.activation = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.activation(x)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        x = self.activation(x + residual)
+        return x
+
+
+class LatentUpsampler(ModelMixin, ConfigMixin):
+    """
+    Model to spatially upsample VAE latents.
+
+    Args:
+        in_channels (`int`): Number of channels in the input latent
+        mid_channels (`int`): Number of channels in the middle layers
+        num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
+        dims (`int`): Number of dimensions for convolutions (2 or 3)
+        spatial_upsample (`bool`): Whether to spatially upsample the latent
+        temporal_upsample (`bool`): Whether to temporally upsample the latent
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 128,
+        mid_channels: int = 512,
+        num_blocks_per_stage: int = 4,
+        dims: int = 3,
+        spatial_upsample: bool = True,
+        temporal_upsample: bool = False,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.mid_channels = mid_channels
+        self.num_blocks_per_stage = num_blocks_per_stage
+        self.dims = dims
+        self.spatial_upsample = spatial_upsample
+        self.temporal_upsample = temporal_upsample
+
+        Conv = nn.Conv2d if dims == 2 else nn.Conv3d
+
+        self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)
+        self.initial_norm = nn.GroupNorm(32, mid_channels)
+        self.initial_activation = nn.SiLU()
+
+        self.res_blocks = nn.ModuleList(
+            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+        )
+
+        if spatial_upsample and temporal_upsample:
+            self.upsampler = nn.Sequential(
+                nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
+                PixelShuffleND(3),
+            )
+        elif spatial_upsample:
+            self.upsampler = nn.Sequential(
+                nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
+                PixelShuffleND(2),
+            )
+        elif temporal_upsample:
+            self.upsampler = nn.Sequential(
+                nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
+                PixelShuffleND(1),
+            )
+        else:
+            raise ValueError(
+                "Either spatial_upsample or temporal_upsample must be True"
+            )
+
+        self.post_upsample_res_blocks = nn.ModuleList(
+            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+        )
+
+        self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)
+
+    def forward(self, latent: torch.Tensor) -> torch.Tensor:
+        b, c, f, h, w = latent.shape
+
+        if self.dims == 2:
+            x = rearrange(latent, "b c f h w -> (b f) c h w")
+            x = self.initial_conv(x)
+            x = self.initial_norm(x)
+            x = self.initial_activation(x)
+
+            for block in self.res_blocks:
+                x = block(x)
+
+            x = self.upsampler(x)
+
+            for block in self.post_upsample_res_blocks:
+                x = block(x)
+
+            x = self.final_conv(x)
+            x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+        else:
+            x = self.initial_conv(latent)
+            x = self.initial_norm(x)
+            x = self.initial_activation(x)
+
+            for block in self.res_blocks:
+                x = block(x)
+
+            if self.temporal_upsample:
+                x = self.upsampler(x)
+                x = x[:, :, 1:, :, :]
+            else:
+                x = rearrange(x, "b c f h w -> (b f) c h w")
+                x = self.upsampler(x)
+                x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+
+            for block in self.post_upsample_res_blocks:
+                x = block(x)
+
+            x = self.final_conv(x)
+
+        return x
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(
+            in_channels=config.get("in_channels", 4),
+            mid_channels=config.get("mid_channels", 128),
+            num_blocks_per_stage=config.get("num_blocks_per_stage", 4),
+            dims=config.get("dims", 2),
+            spatial_upsample=config.get("spatial_upsample", True),
+            temporal_upsample=config.get("temporal_upsample", False),
+        )
+
+    def config(self):
+        return {
+            "_class_name": "LatentUpsampler",
+            "in_channels": self.in_channels,
+            "mid_channels": self.mid_channels,
+            "num_blocks_per_stage": self.num_blocks_per_stage,
+            "dims": self.dims,
+            "spatial_upsample": self.spatial_upsample,
+            "temporal_upsample": self.temporal_upsample,
+        }
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_path: Optional[Union[str, os.PathLike]],
+        *args,
+        **kwargs,
+    ):
+        pretrained_model_path = Path(pretrained_model_path)
+        if pretrained_model_path.is_file() and str(pretrained_model_path).endswith(
+            ".safetensors"
+        ):
+            state_dict = {}
+            with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
+                metadata = f.metadata()
+                for k in f.keys():
+                    state_dict[k] = f.get_tensor(k)
+            config = json.loads(metadata["config"])
+            with torch.device("meta"):
+                latent_upsampler = LatentUpsampler.from_config(config)
+            latent_upsampler.load_state_dict(state_dict, assign=True)
+        return latent_upsampler
+
+
+if __name__ == "__main__":
+    latent_upsampler = LatentUpsampler(num_blocks_per_stage=4, dims=3)
+    print(latent_upsampler)
+    total_params = sum(p.numel() for p in latent_upsampler.parameters())
+    print(f"Total number of parameters: {total_params:,}")
+    latent = torch.randn(1, 128, 9, 16, 16)
+    upsampled_latent = latent_upsampler(latent)
+    print(f"Upsampled latent shape: {upsampled_latent.shape}")

ltx_video/models/autoencoders/pixel_norm.py

@@ -0,0 +1,12 @@
+import torch
+from torch import nn
+
+
+class PixelNorm(nn.Module):
+    def __init__(self, dim=1, eps=1e-8):
+        super(PixelNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x):
+        return x / torch.sqrt(torch.mean(x**2, dim=self.dim, keepdim=True) + self.eps)

ltx_video/models/autoencoders/pixel_shuffle.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+from einops import rearrange
+
+
+class PixelShuffleND(nn.Module):
+    def __init__(self, dims, upscale_factors=(2, 2, 2)):
+        super().__init__()
+        assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
+        self.dims = dims
+        self.upscale_factors = upscale_factors
+
+    def forward(self, x):
+        if self.dims == 3:
+            return rearrange(
+                x,
+                "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+                p1=self.upscale_factors[0],
+                p2=self.upscale_factors[1],
+                p3=self.upscale_factors[2],
+            )
+        elif self.dims == 2:
+            return rearrange(
+                x,
+                "b (c p1 p2) h w -> b c (h p1) (w p2)",
+                p1=self.upscale_factors[0],
+                p2=self.upscale_factors[1],
+            )
+        elif self.dims == 1:
+            return rearrange(
+                x,
+                "b (c p1) f h w -> b c (f p1) h w",
+                p1=self.upscale_factors[0],
+            )

ltx_video/models/autoencoders/vae.py

@@ -0,0 +1,443 @@
+from typing import Optional, Union
+
+import torch
+import inspect
+import math
+import torch.nn as nn
+from diffusers import ConfigMixin, ModelMixin
+from diffusers.models.autoencoders.vae import (
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd
+
+
+class AutoencoderKLWrapper(ModelMixin, ConfigMixin):
+    """Variational Autoencoder (VAE) model with KL loss.
+
+    VAE from the paper Auto-Encoding Variational Bayes by Diederik P. Kingma and Max Welling.
+    This model is a wrapper around an encoder and a decoder, and it adds a KL loss term to the reconstruction loss.
+
+    Args:
+        encoder (`nn.Module`):
+            Encoder module.
+        decoder (`nn.Module`):
+            Decoder module.
+        latent_channels (`int`, *optional*, defaults to 4):
+            Number of latent channels.
+    """
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        latent_channels: int = 4,
+        dims: int = 2,
+        sample_size=512,
+        use_quant_conv: bool = True,
+        normalize_latent_channels: bool = False,
+    ):
+        super().__init__()
+
+
+        self.per_channel_statistics = nn.Module()
+        std_of_means = torch.zeros( (128,), dtype= torch.bfloat16)
+
+        self.per_channel_statistics.register_buffer("std-of-means", std_of_means)
+        self.per_channel_statistics.register_buffer(
+            "mean-of-means",
+                torch.zeros_like(std_of_means)
+        )
+        
+
+
+        # pass init params to Encoder
+        self.encoder = encoder
+        self.use_quant_conv = use_quant_conv
+        self.normalize_latent_channels = normalize_latent_channels
+
+        # pass init params to Decoder
+        quant_dims = 2 if dims == 2 else 3
+        self.decoder = decoder
+        if use_quant_conv:
+            self.quant_conv = make_conv_nd(
+                quant_dims, 2 * latent_channels, 2 * latent_channels, 1
+            )
+            self.post_quant_conv = make_conv_nd(
+                quant_dims, latent_channels, latent_channels, 1
+            )
+        else:
+            self.quant_conv = nn.Identity()
+            self.post_quant_conv = nn.Identity()
+
+        if normalize_latent_channels:
+            if dims == 2:
+                self.latent_norm_out = nn.BatchNorm2d(latent_channels, affine=False)
+            else:
+                self.latent_norm_out = nn.BatchNorm3d(latent_channels, affine=False)
+        else:
+            self.latent_norm_out = nn.Identity()
+        self.use_z_tiling = False
+        self.use_hw_tiling = False
+        self.dims = dims
+        self.z_sample_size = 1
+
+        self.decoder_params = inspect.signature(self.decoder.forward).parameters
+
+        # only relevant if vae tiling is enabled
+        self.set_tiling_params(sample_size=sample_size, overlap_factor=0.25)
+
+    @staticmethod
+    def get_VAE_tile_size(vae_config, device_mem_capacity, mixed_precision):
+
+        z_tile = 4
+        # VAE Tiling
+        if vae_config == 0:
+            if mixed_precision:
+                device_mem_capacity = device_mem_capacity / 1.5
+            if device_mem_capacity >= 24000:
+                use_vae_config = 1            
+            elif device_mem_capacity >= 8000:
+                use_vae_config = 2
+            else:          
+                use_vae_config = 3
+        else:
+            use_vae_config = vae_config
+
+        if use_vae_config == 1:
+            hw_tile = 0 
+        elif use_vae_config == 2:
+            hw_tile = 512  
+        else: 
+            hw_tile = 256  
+
+        return  (z_tile, hw_tile)
+
+    def set_tiling_params(self, sample_size: int = 512, overlap_factor: float = 0.25):
+        self.tile_sample_min_size = sample_size
+        num_blocks = len(self.encoder.down_blocks)
+        # self.tile_latent_min_size = int(sample_size / (2 ** (num_blocks - 1)))
+        self.tile_latent_min_size = int(sample_size / 32)
+        self.tile_overlap_factor = overlap_factor
+
+    def enable_z_tiling(self, z_sample_size: int = 4):
+        r"""
+        Enable tiling during VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor in tiles to compute decoding in several
+        steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_z_tiling = z_sample_size > 1
+        self.z_sample_size = z_sample_size
+        assert (
+            z_sample_size % 4 == 0 or z_sample_size == 1
+        ), f"z_sample_size must be a multiple of 4 or 1. Got {z_sample_size}."
+
+    def disable_z_tiling(self):
+        r"""
+        Disable tiling during VAE decoding. If `use_tiling` was previously invoked, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_z_tiling = False
+
+    def enable_hw_tiling(self):
+        r"""
+        Enable tiling during VAE decoding along the height and width dimension.
+        """
+        self.use_hw_tiling = True
+
+    def disable_hw_tiling(self):
+        r"""
+        Disable tiling during VAE decoding along the height and width dimension.
+        """
+        self.use_hw_tiling = False
+
+    def _hw_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True):
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[3], overlap_size):
+            row = []
+            for j in range(0, x.shape[4], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        moments = torch.cat(result_rows, dim=3)
+        return moments
+
+    def blend_z(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for z in range(blend_extent):
+            b[:, :, z, :, :] = a[:, :, -blend_extent + z, :, :] * (
+                1 - z / blend_extent
+            ) + b[:, :, z, :, :] * (z / blend_extent)
+        return b
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], 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[4], b.shape[4], 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 _hw_tiled_decode(self, z: torch.FloatTensor, target_shape, timestep = None):
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+        tile_target_shape = (
+            *target_shape[:3],
+            self.tile_sample_min_size,
+            self.tile_sample_min_size,
+        )
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[3], overlap_size):
+            row = []
+            for j in range(0, z.shape[4], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile, target_shape=tile_target_shape, timestep = timestep)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        dec = torch.cat(result_rows, dim=3)
+        return dec
+
+    def encode(
+        self, z: torch.FloatTensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
+            tile_latent_min_tsize = self.z_sample_size 
+            tile_sample_min_tsize = tile_latent_min_tsize *  8 
+            tile_overlap_factor = 0.25
+
+            B, C, T, H, W = z.shape
+            overlap_size = int(tile_sample_min_tsize * (1 - tile_overlap_factor))
+            blend_extent = int(tile_latent_min_tsize * tile_overlap_factor)
+            t_limit = tile_latent_min_tsize - blend_extent
+
+            row = []
+            for i in range(0, T, overlap_size):
+                tile = z[:, :, i: i + tile_sample_min_tsize + 1, :, :]
+                if self.use_hw_tiling:
+                    tile = self._hw_tiled_encode(tile, return_dict)
+                else:
+                    tile = self._encode(tile)
+                if i > 0:
+                    tile = tile[:, :, 1:, :, :]
+                row.append(tile)
+            result_row = []
+            for i, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_z(row[i - 1], tile, blend_extent)
+                    result_row.append(tile[:, :, :t_limit, :, :])
+                else:
+                    result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+            moments = torch.cat(result_row, dim=2)
+
+
+        else:
+            moments = (
+                self._hw_tiled_encode(z, return_dict)
+                if self.use_hw_tiling and z.shape[2] > 1 
+                else self._encode(z)
+            )
+
+        posterior = DiagonalGaussianDistribution(moments)
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _normalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        if isinstance(self.latent_norm_out, nn.BatchNorm3d):
+            _, c, _, _, _ = z.shape
+            z = torch.cat(
+                [
+                    self.latent_norm_out(z[:, : c // 2, :, :, :]),
+                    z[:, c // 2 :, :, :, :],
+                ],
+                dim=1,
+            )
+        elif isinstance(self.latent_norm_out, nn.BatchNorm2d):
+            raise NotImplementedError("BatchNorm2d not supported")
+        return z
+
+    def _unnormalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        if isinstance(self.latent_norm_out, nn.BatchNorm3d):
+            running_mean = self.latent_norm_out.running_mean.view(1, -1, 1, 1, 1)
+            running_var = self.latent_norm_out.running_var.view(1, -1, 1, 1, 1)
+            eps = self.latent_norm_out.eps
+
+            z = z * torch.sqrt(running_var + eps) + running_mean
+        elif isinstance(self.latent_norm_out, nn.BatchNorm3d):
+            raise NotImplementedError("BatchNorm2d not supported")
+        return z
+
+    def _encode(self, x: torch.FloatTensor) -> AutoencoderKLOutput:
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        moments = self._normalize_latent_channels(moments)
+        return moments
+
+    def _decode(
+        self,
+        z: torch.FloatTensor,
+        target_shape=None,
+        timestep: Optional[torch.Tensor] = None,
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        z = self._unnormalize_latent_channels(z)
+        z = self.post_quant_conv(z)
+        if "timestep" in self.decoder_params:
+            dec = self.decoder(z, target_shape=target_shape, timestep=timestep)
+        else:
+            dec = self.decoder(z, target_shape=target_shape)
+        return dec
+
+    def decode(
+        self,
+        z: torch.FloatTensor,
+        return_dict: bool = True,
+        target_shape=None,
+        timestep: Optional[torch.Tensor] = None,
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        assert target_shape is not None, "target_shape must be provided for decoding"
+        if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
+        # Split z into overlapping tiles and decode them separately.
+            tile_latent_min_tsize = self.z_sample_size 
+            tile_sample_min_tsize = tile_latent_min_tsize *  8 
+            tile_overlap_factor = 0.25
+
+            B, C, T, H, W = z.shape
+            overlap_size = int(tile_latent_min_tsize * (1 - tile_overlap_factor))
+            blend_extent = int(tile_sample_min_tsize * tile_overlap_factor)
+            t_limit = tile_sample_min_tsize - blend_extent
+
+            row = []
+            for i in range(0, T, overlap_size):
+                tile = z[:, :, i: i + tile_latent_min_tsize + 1, :, :]
+                target_shape_split = list(target_shape)
+                target_shape_split[2] = tile.shape[2] * 8                
+                if self.use_hw_tiling:
+                    decoded = self._hw_tiled_decode(tile, target_shape, timestep)
+                else:
+                    decoded = self._decode(tile, target_shape=target_shape, timestep=timestep)
+
+                if i > 0:
+                    decoded = decoded[:, :, 1:, :, :]
+                row.append(decoded.to(torch.float16).cpu())
+                decoded = None
+            result_row = []
+            for i, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_z(row[i - 1], tile, blend_extent)
+                    result_row.append(tile[:, :, :t_limit, :, :])
+                else:
+                    result_row.append(tile[:, :, :t_limit + 1, :, :])
+
+            dec = torch.cat(result_row, dim=2)
+            if not return_dict:
+                return (dec,)
+
+            return DecoderOutput(sample=dec)            
+        else:
+            decoded = (
+                self._hw_tiled_decode(z, target_shape, timestep)
+                if self.use_hw_tiling
+                else self._decode(z, target_shape=target_shape, timestep=timestep)
+            )
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                Generator used to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z, target_shape=sample.shape).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)

ltx_video/models/autoencoders/vae_encode.py

@@ -0,0 +1,247 @@
+from typing import Tuple
+import torch
+from diffusers import AutoencoderKL
+from einops import rearrange
+from torch import Tensor
+
+
+from ltx_video.models.autoencoders.causal_video_autoencoder import (
+    CausalVideoAutoencoder,
+)
+from ltx_video.models.autoencoders.video_autoencoder import (
+    Downsample3D,
+    VideoAutoencoder,
+)
+
+try:
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    xm = None
+
+
+def vae_encode(
+    media_items: Tensor,
+    vae: AutoencoderKL,
+    split_size: int = 1,
+    vae_per_channel_normalize=False,
+) -> Tensor:
+    """
+    Encodes media items (images or videos) into latent representations using a specified VAE model.
+    The function supports processing batches of images or video frames and can handle the processing
+    in smaller sub-batches if needed.
+
+    Args:
+        media_items (Tensor): A torch Tensor containing the media items to encode. The expected
+            shape is (batch_size, channels, height, width) for images or (batch_size, channels,
+            frames, height, width) for videos.
+        vae (AutoencoderKL): An instance of the `AutoencoderKL` class from the `diffusers` library,
+            pre-configured and loaded with the appropriate model weights.
+        split_size (int, optional): The number of sub-batches to split the input batch into for encoding.
+            If set to more than 1, the input media items are processed in smaller batches according to
+            this value. Defaults to 1, which processes all items in a single batch.
+
+    Returns:
+        Tensor: A torch Tensor of the encoded latent representations. The shape of the tensor is adjusted
+            to match the input shape, scaled by the model's configuration.
+
+    Examples:
+        >>> import torch
+        >>> from diffusers import AutoencoderKL
+        >>> vae = AutoencoderKL.from_pretrained('your-model-name')
+        >>> images = torch.rand(10, 3, 8 256, 256)  # Example tensor with 10 videos of 8 frames.
+        >>> latents = vae_encode(images, vae)
+        >>> print(latents.shape)  # Output shape will depend on the model's latent configuration.
+
+    Note:
+        In case of a video, the function encodes the media item frame-by frame.
+    """
+    is_video_shaped = media_items.dim() == 5
+    batch_size, channels = media_items.shape[0:2]
+
+    if channels != 3:
+        raise ValueError(f"Expects tensors with 3 channels, got {channels}.")
+
+    if is_video_shaped and not isinstance(
+        vae, (VideoAutoencoder, CausalVideoAutoencoder)
+    ):
+        media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
+    if split_size > 1:
+        if len(media_items) % split_size != 0:
+            raise ValueError(
+                "Error: The batch size must be divisible by 'train.vae_bs_split"
+            )
+        encode_bs = len(media_items) // split_size
+        # latents = [vae.encode(image_batch).latent_dist.sample() for image_batch in media_items.split(encode_bs)]
+        latents = []
+        if media_items.device.type == "xla":
+            xm.mark_step()
+        for image_batch in media_items.split(encode_bs):
+            latents.append(vae.encode(image_batch).latent_dist.sample())
+            if media_items.device.type == "xla":
+                xm.mark_step()
+        latents = torch.cat(latents, dim=0)
+    else:
+        latents = vae.encode(media_items).latent_dist.sample()
+
+    latents = normalize_latents(latents, vae, vae_per_channel_normalize)
+    if is_video_shaped and not isinstance(
+        vae, (VideoAutoencoder, CausalVideoAutoencoder)
+    ):
+        latents = rearrange(latents, "(b n) c h w -> b c n h w", b=batch_size)
+    return latents
+
+
+def vae_decode(
+    latents: Tensor,
+    vae: AutoencoderKL,
+    is_video: bool = True,
+    split_size: int = 1,
+    vae_per_channel_normalize=False,
+    timestep=None,
+) -> Tensor:
+    is_video_shaped = latents.dim() == 5
+    batch_size = latents.shape[0]
+
+    if is_video_shaped and not isinstance(
+        vae, (VideoAutoencoder, CausalVideoAutoencoder)
+    ):
+        latents = rearrange(latents, "b c n h w -> (b n) c h w")
+    if split_size > 1:
+        if len(latents) % split_size != 0:
+            raise ValueError(
+                "Error: The batch size must be divisible by 'train.vae_bs_split"
+            )
+        encode_bs = len(latents) // split_size
+        image_batch = [
+            _run_decoder(
+                latent_batch, vae, is_video, vae_per_channel_normalize, timestep
+            )
+            for latent_batch in latents.split(encode_bs)
+        ]
+        images = torch.cat(image_batch, dim=0)
+    else:
+        images = _run_decoder(
+            latents, vae, is_video, vae_per_channel_normalize, timestep
+        )
+
+    if is_video_shaped and not isinstance(
+        vae, (VideoAutoencoder, CausalVideoAutoencoder)
+    ):
+        images = rearrange(images, "(b n) c h w -> b c n h w", b=batch_size)
+    return images
+
+
+def _run_decoder(
+    latents: Tensor,
+    vae: AutoencoderKL,
+    is_video: bool,
+    vae_per_channel_normalize=False,
+    timestep=None,
+) -> Tensor:
+    if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)):
+        *_, fl, hl, wl = latents.shape
+        temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae)
+        latents = latents.to(vae.dtype)
+        vae_decode_kwargs = {}
+        if timestep is not None:
+            vae_decode_kwargs["timestep"] = timestep
+        image = vae.decode(
+            un_normalize_latents(latents, vae, vae_per_channel_normalize),
+            return_dict=False,
+            target_shape=(
+                1,
+                3,
+                fl * temporal_scale if is_video else 1,
+                hl * spatial_scale,
+                wl * spatial_scale,
+            ),
+            **vae_decode_kwargs,
+        )[0]
+    else:
+        image = vae.decode(
+            un_normalize_latents(latents, vae, vae_per_channel_normalize),
+            return_dict=False,
+        )[0]
+    return image
+
+
+def get_vae_size_scale_factor(vae: AutoencoderKL) -> float:
+    if isinstance(vae, CausalVideoAutoencoder):
+        spatial = vae.spatial_downscale_factor
+        temporal = vae.temporal_downscale_factor
+    else:
+        down_blocks = len(
+            [
+                block
+                for block in vae.encoder.down_blocks
+                if isinstance(block.downsample, Downsample3D)
+            ]
+        )
+        spatial = vae.config.patch_size * 2**down_blocks
+        temporal = (
+            vae.config.patch_size_t * 2**down_blocks
+            if isinstance(vae, VideoAutoencoder)
+            else 1
+        )
+
+    return (temporal, spatial, spatial)
+
+
+def latent_to_pixel_coords(
+    latent_coords: Tensor, vae: AutoencoderKL, causal_fix: bool = False
+) -> Tensor:
+    """
+    Converts latent coordinates to pixel coordinates by scaling them according to the VAE's
+    configuration.
+
+    Args:
+        latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents]
+        containing the latent corner coordinates of each token.
+        vae (AutoencoderKL): The VAE model
+        causal_fix (bool): Whether to take into account the different temporal scale
+            of the first frame. Default = False for backwards compatibility.
+    Returns:
+        Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates.
+    """
+
+    scale_factors = get_vae_size_scale_factor(vae)
+    causal_fix = isinstance(vae, CausalVideoAutoencoder) and causal_fix
+    pixel_coords = latent_to_pixel_coords_from_factors(
+        latent_coords, scale_factors, causal_fix
+    )
+    return pixel_coords
+
+
+def latent_to_pixel_coords_from_factors(
+    latent_coords: Tensor, scale_factors: Tuple, causal_fix: bool = False
+) -> Tensor:
+    pixel_coords = (
+        latent_coords
+        * torch.tensor(scale_factors, device=latent_coords.device)[None, :, None]
+    )
+    if causal_fix:
+        # Fix temporal scale for first frame to 1 due to causality
+        pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0)
+    return pixel_coords
+
+
+def normalize_latents(
+    latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
+) -> Tensor:
+    return (
+        (latents - vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1))
+        / vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
+        if vae_per_channel_normalize
+        else latents * vae.config.scaling_factor
+    )
+
+
+def un_normalize_latents(
+    latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
+) -> Tensor:
+    return (
+        latents * vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
+        + vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
+        if vae_per_channel_normalize
+        else latents / vae.config.scaling_factor
+    )