diffusers 0.32.1__py3-none-any.whl → 0.33.0__py3-none-any.whl

diffusers/models/autoencoders/autoencoder_kl_magvit.py

@@ -0,0 +1,1094 @@
+# Copyright 2025 The EasyAnimate team and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...utils import logging
+from ...utils.accelerate_utils import apply_forward_hook
+from ..activations import get_activation
+from ..modeling_outputs import AutoencoderKLOutput
+from ..modeling_utils import ModelMixin
+from .vae import DecoderOutput, DiagonalGaussianDistribution
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class EasyAnimateCausalConv3d(nn.Conv3d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, ...]] = 3,
+        stride: Union[int, Tuple[int, ...]] = 1,
+        padding: Union[int, Tuple[int, ...]] = 1,
+        dilation: Union[int, Tuple[int, ...]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+    ):
+        # Ensure kernel_size, stride, and dilation are tuples of length 3
+        kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3
+        assert len(kernel_size) == 3, f"Kernel size must be a 3-tuple, got {kernel_size} instead."
+
+        stride = stride if isinstance(stride, tuple) else (stride,) * 3
+        assert len(stride) == 3, f"Stride must be a 3-tuple, got {stride} instead."
+
+        dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3
+        assert len(dilation) == 3, f"Dilation must be a 3-tuple, got {dilation} instead."
+
+        # Unpack kernel size, stride, and dilation for temporal, height, and width dimensions
+        t_ks, h_ks, w_ks = kernel_size
+        self.t_stride, h_stride, w_stride = stride
+        t_dilation, h_dilation, w_dilation = dilation
+
+        # Calculate padding for temporal dimension to maintain causality
+        t_pad = (t_ks - 1) * t_dilation
+
+        # Calculate padding for height and width dimensions based on the padding parameter
+        if padding is None:
+            h_pad = math.ceil(((h_ks - 1) * h_dilation + (1 - h_stride)) / 2)
+            w_pad = math.ceil(((w_ks - 1) * w_dilation + (1 - w_stride)) / 2)
+        elif isinstance(padding, int):
+            h_pad = w_pad = padding
+        else:
+            assert NotImplementedError
+
+        # Store temporal padding and initialize flags and previous features cache
+        self.temporal_padding = t_pad
+        self.temporal_padding_origin = math.ceil(((t_ks - 1) * w_dilation + (1 - w_stride)) / 2)
+
+        self.prev_features = None
+
+        # Initialize the parent class with modified padding
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            padding=(0, h_pad, w_pad),
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
+
+    def _clear_conv_cache(self):
+        del self.prev_features
+        self.prev_features = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Ensure input tensor is of the correct type
+        dtype = hidden_states.dtype
+        if self.prev_features is None:
+            # Pad the input tensor in the temporal dimension to maintain causality
+            hidden_states = F.pad(
+                hidden_states,
+                pad=(0, 0, 0, 0, self.temporal_padding, 0),
+                mode="replicate",  # TODO: check if this is necessary
+            )
+            hidden_states = hidden_states.to(dtype=dtype)
+
+            # Clear cache before processing and store previous features for causality
+            self._clear_conv_cache()
+            self.prev_features = hidden_states[:, :, -self.temporal_padding :].clone()
+
+            # Process the input tensor in chunks along the temporal dimension
+            num_frames = hidden_states.size(2)
+            outputs = []
+            i = 0
+            while i + self.temporal_padding + 1 <= num_frames:
+                out = super().forward(hidden_states[:, :, i : i + self.temporal_padding + 1])
+                i += self.t_stride
+                outputs.append(out)
+            return torch.concat(outputs, 2)
+        else:
+            # Concatenate previous features with the input tensor for continuous temporal processing
+            if self.t_stride == 2:
+                hidden_states = torch.concat(
+                    [self.prev_features[:, :, -(self.temporal_padding - 1) :], hidden_states], dim=2
+                )
+            else:
+                hidden_states = torch.concat([self.prev_features, hidden_states], dim=2)
+            hidden_states = hidden_states.to(dtype=dtype)
+
+            # Clear cache and update previous features
+            self._clear_conv_cache()
+            self.prev_features = hidden_states[:, :, -self.temporal_padding :].clone()
+
+            # Process the concatenated tensor in chunks along the temporal dimension
+            num_frames = hidden_states.size(2)
+            outputs = []
+            i = 0
+            while i + self.temporal_padding + 1 <= num_frames:
+                out = super().forward(hidden_states[:, :, i : i + self.temporal_padding + 1])
+                i += self.t_stride
+                outputs.append(out)
+            return torch.concat(outputs, 2)
+
+
+class EasyAnimateResidualBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        non_linearity: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_group_norm: bool = True,
+        dropout: float = 0.0,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+
+        self.output_scale_factor = output_scale_factor
+
+        # Group normalization for input tensor
+        self.norm1 = nn.GroupNorm(
+            num_groups=norm_num_groups,
+            num_channels=in_channels,
+            eps=norm_eps,
+            affine=True,
+        )
+        self.nonlinearity = get_activation(non_linearity)
+        self.conv1 = EasyAnimateCausalConv3d(in_channels, out_channels, kernel_size=3)
+
+        self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=norm_eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = EasyAnimateCausalConv3d(out_channels, out_channels, kernel_size=3)
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+        else:
+            self.shortcut = nn.Identity()
+
+        self.spatial_group_norm = spatial_group_norm
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        shortcut = self.shortcut(hidden_states)
+
+        if self.spatial_group_norm:
+            batch_size = hidden_states.size(0)
+            hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)  # [B, C, T, H, W] -> [B * T, C, H, W]
+            hidden_states = self.norm1(hidden_states)
+            hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(
+                0, 2, 1, 3, 4
+            )  # [B * T, C, H, W] -> [B, C, T, H, W]
+        else:
+            hidden_states = self.norm1(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        if self.spatial_group_norm:
+            batch_size = hidden_states.size(0)
+            hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)  # [B, C, T, H, W] -> [B * T, C, H, W]
+            hidden_states = self.norm2(hidden_states)
+            hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(
+                0, 2, 1, 3, 4
+            )  # [B * T, C, H, W] -> [B, C, T, H, W]
+        else:
+            hidden_states = self.norm2(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        return (hidden_states + shortcut) / self.output_scale_factor
+
+
+class EasyAnimateDownsampler3D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: tuple = (2, 2, 2)):
+        super().__init__()
+
+        self.conv = EasyAnimateCausalConv3d(
+            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=0
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.pad(hidden_states, (0, 1, 0, 1))
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class EasyAnimateUpsampler3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        temporal_upsample: bool = False,
+        spatial_group_norm: bool = True,
+    ):
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.temporal_upsample = temporal_upsample
+        self.spatial_group_norm = spatial_group_norm
+
+        self.conv = EasyAnimateCausalConv3d(
+            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size
+        )
+        self.prev_features = None
+
+    def _clear_conv_cache(self):
+        del self.prev_features
+        self.prev_features = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest")
+        hidden_states = self.conv(hidden_states)
+
+        if self.temporal_upsample:
+            if self.prev_features is None:
+                self.prev_features = hidden_states
+            else:
+                hidden_states = F.interpolate(
+                    hidden_states,
+                    scale_factor=(2, 1, 1),
+                    mode="trilinear" if not self.spatial_group_norm else "nearest",
+                )
+        return hidden_states
+
+
+class EasyAnimateDownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_group_norm: bool = True,
+        dropout: float = 0.0,
+        output_scale_factor: float = 1.0,
+        add_downsample: bool = True,
+        add_temporal_downsample: bool = True,
+    ):
+        super().__init__()
+
+        self.convs = nn.ModuleList([])
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            self.convs.append(
+                EasyAnimateResidualBlock3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    non_linearity=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=norm_eps,
+                    spatial_group_norm=spatial_group_norm,
+                    dropout=dropout,
+                    output_scale_factor=output_scale_factor,
+                )
+            )
+
+        if add_downsample and add_temporal_downsample:
+            self.downsampler = EasyAnimateDownsampler3D(out_channels, out_channels, kernel_size=3, stride=(2, 2, 2))
+            self.spatial_downsample_factor = 2
+            self.temporal_downsample_factor = 2
+        elif add_downsample and not add_temporal_downsample:
+            self.downsampler = EasyAnimateDownsampler3D(out_channels, out_channels, kernel_size=3, stride=(1, 2, 2))
+            self.spatial_downsample_factor = 2
+            self.temporal_downsample_factor = 1
+        else:
+            self.downsampler = None
+            self.spatial_downsample_factor = 1
+            self.temporal_downsample_factor = 1
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for conv in self.convs:
+            hidden_states = conv(hidden_states)
+        if self.downsampler is not None:
+            hidden_states = self.downsampler(hidden_states)
+        return hidden_states
+
+
+class EasyAnimateUpBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_group_norm: bool = False,
+        dropout: float = 0.0,
+        output_scale_factor: float = 1.0,
+        add_upsample: bool = True,
+        add_temporal_upsample: bool = True,
+    ):
+        super().__init__()
+
+        self.convs = nn.ModuleList([])
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            self.convs.append(
+                EasyAnimateResidualBlock3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    non_linearity=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=norm_eps,
+                    spatial_group_norm=spatial_group_norm,
+                    dropout=dropout,
+                    output_scale_factor=output_scale_factor,
+                )
+            )
+
+        if add_upsample:
+            self.upsampler = EasyAnimateUpsampler3D(
+                in_channels,
+                in_channels,
+                temporal_upsample=add_temporal_upsample,
+                spatial_group_norm=spatial_group_norm,
+            )
+        else:
+            self.upsampler = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for conv in self.convs:
+            hidden_states = conv(hidden_states)
+        if self.upsampler is not None:
+            hidden_states = self.upsampler(hidden_states)
+        return hidden_states
+
+
+class EasyAnimateMidBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_layers: int = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_group_norm: bool = True,
+        dropout: float = 0.0,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+
+        norm_num_groups = norm_num_groups if norm_num_groups is not None else min(in_channels // 4, 32)
+
+        self.convs = nn.ModuleList(
+            [
+                EasyAnimateResidualBlock3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    non_linearity=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=norm_eps,
+                    spatial_group_norm=spatial_group_norm,
+                    dropout=dropout,
+                    output_scale_factor=output_scale_factor,
+                )
+            ]
+        )
+
+        for _ in range(num_layers - 1):
+            self.convs.append(
+                EasyAnimateResidualBlock3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    non_linearity=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=norm_eps,
+                    spatial_group_norm=spatial_group_norm,
+                    dropout=dropout,
+                    output_scale_factor=output_scale_factor,
+                )
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.convs[0](hidden_states)
+        for resnet in self.convs[1:]:
+            hidden_states = resnet(hidden_states)
+        return hidden_states
+
+
+class EasyAnimateEncoder(nn.Module):
+    r"""
+    Causal encoder for 3D video-like data used in [EasyAnimate](https://arxiv.org/abs/2405.18991).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 8,
+        down_block_types: Tuple[str, ...] = (
+            "SpatialDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        spatial_group_norm: bool = False,
+    ):
+        super().__init__()
+
+        # 1. Input convolution
+        self.conv_in = EasyAnimateCausalConv3d(in_channels, block_out_channels[0], kernel_size=3)
+
+        # 2. Down blocks
+        self.down_blocks = nn.ModuleList([])
+        output_channels = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channels = output_channels
+            output_channels = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            if down_block_type == "SpatialDownBlock3D":
+                down_block = EasyAnimateDownBlock3D(
+                    in_channels=input_channels,
+                    out_channels=output_channels,
+                    num_layers=layers_per_block,
+                    act_fn=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=1e-6,
+                    spatial_group_norm=spatial_group_norm,
+                    add_downsample=not is_final_block,
+                    add_temporal_downsample=False,
+                )
+            elif down_block_type == "SpatialTemporalDownBlock3D":
+                down_block = EasyAnimateDownBlock3D(
+                    in_channels=input_channels,
+                    out_channels=output_channels,
+                    num_layers=layers_per_block,
+                    act_fn=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=1e-6,
+                    spatial_group_norm=spatial_group_norm,
+                    add_downsample=not is_final_block,
+                    add_temporal_downsample=True,
+                )
+            else:
+                raise ValueError(f"Unknown up block type: {down_block_type}")
+            self.down_blocks.append(down_block)
+
+        # 3. Middle block
+        self.mid_block = EasyAnimateMidBlock3d(
+            in_channels=block_out_channels[-1],
+            num_layers=layers_per_block,
+            act_fn=act_fn,
+            spatial_group_norm=spatial_group_norm,
+            norm_num_groups=norm_num_groups,
+            norm_eps=1e-6,
+            dropout=0,
+            output_scale_factor=1,
+        )
+
+        # 4. Output normalization & convolution
+        self.spatial_group_norm = spatial_group_norm
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[-1],
+            num_groups=norm_num_groups,
+            eps=1e-6,
+        )
+        self.conv_act = get_activation(act_fn)
+
+        # Initialize the output convolution layer
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = EasyAnimateCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states: (B, C, T, H, W)
+        hidden_states = self.conv_in(hidden_states)
+
+        for down_block in self.down_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(down_block, hidden_states)
+            else:
+                hidden_states = down_block(hidden_states)
+
+        hidden_states = self.mid_block(hidden_states)
+
+        if self.spatial_group_norm:
+            batch_size = hidden_states.size(0)
+            hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+            hidden_states = self.conv_norm_out(hidden_states)
+            hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        else:
+            hidden_states = self.conv_norm_out(hidden_states)
+
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+        return hidden_states
+
+
+class EasyAnimateDecoder(nn.Module):
+    r"""
+    Causal decoder for 3D video-like data used in [EasyAnimate](https://arxiv.org/abs/2405.18991).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        in_channels: int = 8,
+        out_channels: int = 3,
+        up_block_types: Tuple[str, ...] = (
+            "SpatialUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        spatial_group_norm: bool = False,
+    ):
+        super().__init__()
+
+        # 1. Input convolution
+        self.conv_in = EasyAnimateCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3)
+
+        # 2. Middle block
+        self.mid_block = EasyAnimateMidBlock3d(
+            in_channels=block_out_channels[-1],
+            num_layers=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            norm_eps=1e-6,
+            dropout=0,
+            output_scale_factor=1,
+        )
+
+        # 3. Up blocks
+        self.up_blocks = nn.ModuleList([])
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channels = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            input_channels = output_channels
+            output_channels = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            # Create and append up block to up_blocks
+            if up_block_type == "SpatialUpBlock3D":
+                up_block = EasyAnimateUpBlock3d(
+                    in_channels=input_channels,
+                    out_channels=output_channels,
+                    num_layers=layers_per_block + 1,
+                    act_fn=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=1e-6,
+                    spatial_group_norm=spatial_group_norm,
+                    add_upsample=not is_final_block,
+                    add_temporal_upsample=False,
+                )
+            elif up_block_type == "SpatialTemporalUpBlock3D":
+                up_block = EasyAnimateUpBlock3d(
+                    in_channels=input_channels,
+                    out_channels=output_channels,
+                    num_layers=layers_per_block + 1,
+                    act_fn=act_fn,
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=1e-6,
+                    spatial_group_norm=spatial_group_norm,
+                    add_upsample=not is_final_block,
+                    add_temporal_upsample=True,
+                )
+            else:
+                raise ValueError(f"Unknown up block type: {up_block_type}")
+            self.up_blocks.append(up_block)
+
+        # Output normalization and activation
+        self.spatial_group_norm = spatial_group_norm
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[0],
+            num_groups=norm_num_groups,
+            eps=1e-6,
+        )
+        self.conv_act = get_activation(act_fn)
+
+        # Output convolution layer
+        self.conv_out = EasyAnimateCausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states: (B, C, T, H, W)
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
+        else:
+            hidden_states = self.mid_block(hidden_states)
+
+        for up_block in self.up_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(up_block, hidden_states)
+            else:
+                hidden_states = up_block(hidden_states)
+
+        if self.spatial_group_norm:
+            batch_size = hidden_states.size(0)
+            hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)  # [B, C, T, H, W] -> [B * T, C, H, W]
+            hidden_states = self.conv_norm_out(hidden_states)
+            hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(
+                0, 2, 1, 3, 4
+            )  # [B * T, C, H, W] -> [B, C, T, H, W]
+        else:
+            hidden_states = self.conv_norm_out(hidden_states)
+
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+        return hidden_states
+
+
+class AutoencoderKLMagvit(ModelMixin, ConfigMixin):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This
+    model is used in [EasyAnimate](https://arxiv.org/abs/2405.18991).
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        latent_channels: int = 16,
+        out_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = [128, 256, 512, 512],
+        down_block_types: Tuple[str, ...] = [
+            "SpatialDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+            "SpatialTemporalDownBlock3D",
+        ],
+        up_block_types: Tuple[str, ...] = [
+            "SpatialUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+            "SpatialTemporalUpBlock3D",
+        ],
+        layers_per_block: int = 2,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        scaling_factor: float = 0.7125,
+        spatial_group_norm: bool = True,
+    ):
+        super().__init__()
+
+        # Initialize the encoder
+        self.encoder = EasyAnimateEncoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            double_z=True,
+            spatial_group_norm=spatial_group_norm,
+        )
+
+        # Initialize the decoder
+        self.decoder = EasyAnimateDecoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            spatial_group_norm=spatial_group_norm,
+        )
+
+        # Initialize convolution layers for quantization and post-quantization
+        self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
+        self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
+
+        self.spatial_compression_ratio = 2 ** (len(block_out_channels) - 1)
+        self.temporal_compression_ratio = 2 ** (len(block_out_channels) - 2)
+
+        # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
+        # to perform decoding of a single video latent at a time.
+        self.use_slicing = False
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
+
+        # When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames
+        # at a fixed frame batch size (based on `self.num_latent_frames_batch_size`), the memory requirement can be lowered.
+        self.use_framewise_encoding = False
+        self.use_framewise_decoding = False
+
+        # Assign mini-batch sizes for encoder and decoder
+        self.num_sample_frames_batch_size = 4
+        self.num_latent_frames_batch_size = 1
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 512
+        self.tile_sample_min_width = 512
+        self.tile_sample_min_num_frames = 4
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 448
+        self.tile_sample_stride_width = 448
+        self.tile_sample_stride_num_frames = 8
+
+    def _clear_conv_cache(self):
+        # Clear cache for convolutional layers if needed
+        for name, module in self.named_modules():
+            if isinstance(module, EasyAnimateCausalConv3d):
+                module._clear_conv_cache()
+            if isinstance(module, EasyAnimateUpsampler3D):
+                module._clear_conv_cache()
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: Optional[int] = None,
+        tile_sample_min_width: Optional[int] = None,
+        tile_sample_min_num_frames: Optional[int] = None,
+        tile_sample_stride_height: Optional[float] = None,
+        tile_sample_stride_width: Optional[float] = None,
+        tile_sample_stride_num_frames: Optional[float] = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+        """
+        self.use_tiling = True
+        self.use_framewise_decoding = True
+        self.use_framewise_encoding = True
+        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames
+        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
+        self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
+
+    def disable_tiling(self) -> None:
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_tiling = False
+
+    def enable_slicing(self) -> None:
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self) -> None:
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @apply_forward_hook
+    def _encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_height or x.shape[-2] > self.tile_sample_min_width):
+            return self.tiled_encode(x, return_dict=return_dict)
+
+        first_frames = self.encoder(x[:, :, :1, :, :])
+        h = [first_frames]
+        for i in range(1, x.shape[2], self.num_sample_frames_batch_size):
+            next_frames = self.encoder(x[:, :, i : i + self.num_sample_frames_batch_size, :, :])
+            h.append(next_frames)
+        h = torch.cat(h, dim=2)
+        moments = self.quant_conv(h)
+
+        self._clear_conv_cache()
+        return moments
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded videos. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        if self.use_tiling and (z.shape[-1] > tile_latent_min_height or z.shape[-2] > tile_latent_min_width):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        z = self.post_quant_conv(z)
+
+        # Process the first frame and save the result
+        first_frames = self.decoder(z[:, :, :1, :, :])
+        # Initialize the list to store the processed frames, starting with the first frame
+        dec = [first_frames]
+        # Process the remaining frames, with the number of frames processed at a time determined by mini_batch_decoder
+        for i in range(1, z.shape[2], self.num_latent_frames_batch_size):
+            next_frames = self.decoder(z[:, :, i : i + self.num_latent_frames_batch_size, :, :])
+            dec.append(next_frames)
+        # Concatenate all processed frames along the channel dimension
+        dec = torch.cat(dec, dim=2)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        self._clear_conv_cache()
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)
+
+    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 tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput:
+        batch_size, num_channels, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_height,
+                    j : j + self.tile_sample_min_width,
+                ]
+
+                first_frames = self.encoder(tile[:, :, 0:1, :, :])
+                tile_h = [first_frames]
+                for k in range(1, num_frames, self.num_sample_frames_batch_size):
+                    next_frames = self.encoder(tile[:, :, k : k + self.num_sample_frames_batch_size, :, :])
+                    tile_h.append(next_frames)
+                tile = torch.cat(tile_h, dim=2)
+                tile = self.quant_conv(tile)
+                self._clear_conv_cache()
+                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_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :latent_height, :latent_width])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        moments = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return moments
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                tile = self.post_quant_conv(tile)
+
+                # Process the first frame and save the result
+                first_frames = self.decoder(tile[:, :, :1, :, :])
+                # Initialize the list to store the processed frames, starting with the first frame
+                tile_dec = [first_frames]
+                # Process the remaining frames, with the number of frames processed at a time determined by mini_batch_decoder
+                for k in range(1, num_frames, self.num_latent_frames_batch_size):
+                    next_frames = self.decoder(tile[:, :, k : k + self.num_latent_frames_batch_size, :, :])
+                    tile_dec.append(next_frames)
+                # Concatenate all processed frames along the channel dimension
+                decoded = torch.cat(tile_dec, dim=2)
+                self._clear_conv_cache()
+                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_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)