diffusers 0.27.0__py3-none-any.whl → 0.32.2__py3-none-any.whl

diffusers/models/autoencoders/autoencoder_kl_allegro.py

@@ -0,0 +1,1149 @@
+# Copyright 2024 The RhymesAI 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
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...utils.accelerate_utils import apply_forward_hook
+from ..attention_processor import Attention, SpatialNorm
+from ..autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution
+from ..downsampling import Downsample2D
+from ..modeling_outputs import AutoencoderKLOutput
+from ..modeling_utils import ModelMixin
+from ..resnet import ResnetBlock2D
+from ..upsampling import Upsample2D
+
+
+class AllegroTemporalConvLayer(nn.Module):
+    r"""
+    Temporal convolutional layer that can be used for video (sequence of images) input. Code adapted from:
+    https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        up_sample: bool = False,
+        down_sample: bool = False,
+        stride: int = 1,
+    ) -> None:
+        super().__init__()
+
+        out_dim = out_dim or in_dim
+        pad_h = pad_w = int((stride - 1) * 0.5)
+        pad_t = 0
+
+        self.down_sample = down_sample
+        self.up_sample = up_sample
+
+        if down_sample:
+            self.conv1 = nn.Sequential(
+                nn.GroupNorm(norm_num_groups, in_dim),
+                nn.SiLU(),
+                nn.Conv3d(in_dim, out_dim, (2, stride, stride), stride=(2, 1, 1), padding=(0, pad_h, pad_w)),
+            )
+        elif up_sample:
+            self.conv1 = nn.Sequential(
+                nn.GroupNorm(norm_num_groups, in_dim),
+                nn.SiLU(),
+                nn.Conv3d(in_dim, out_dim * 2, (1, stride, stride), padding=(0, pad_h, pad_w)),
+            )
+        else:
+            self.conv1 = nn.Sequential(
+                nn.GroupNorm(norm_num_groups, in_dim),
+                nn.SiLU(),
+                nn.Conv3d(in_dim, out_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_w)),
+            )
+        self.conv2 = nn.Sequential(
+            nn.GroupNorm(norm_num_groups, out_dim),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_w)),
+        )
+        self.conv3 = nn.Sequential(
+            nn.GroupNorm(norm_num_groups, out_dim),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_h)),
+        )
+        self.conv4 = nn.Sequential(
+            nn.GroupNorm(norm_num_groups, out_dim),
+            nn.SiLU(),
+            nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_h)),
+        )
+
+    @staticmethod
+    def _pad_temporal_dim(hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = torch.cat((hidden_states[:, :, 0:1], hidden_states), dim=2)
+        hidden_states = torch.cat((hidden_states, hidden_states[:, :, -1:]), dim=2)
+        return hidden_states
+
+    def forward(self, hidden_states: torch.Tensor, batch_size: int) -> torch.Tensor:
+        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+
+        if self.down_sample:
+            identity = hidden_states[:, :, ::2]
+        elif self.up_sample:
+            identity = hidden_states.repeat_interleave(2, dim=2)
+        else:
+            identity = hidden_states
+
+        if self.down_sample or self.up_sample:
+            hidden_states = self.conv1(hidden_states)
+        else:
+            hidden_states = self._pad_temporal_dim(hidden_states)
+            hidden_states = self.conv1(hidden_states)
+
+        if self.up_sample:
+            hidden_states = hidden_states.unflatten(1, (2, -1)).permute(0, 2, 3, 1, 4, 5).flatten(2, 3)
+
+        hidden_states = self._pad_temporal_dim(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        hidden_states = self._pad_temporal_dim(hidden_states)
+        hidden_states = self.conv3(hidden_states)
+
+        hidden_states = self._pad_temporal_dim(hidden_states)
+        hidden_states = self.conv4(hidden_states)
+
+        hidden_states = identity + hidden_states
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+
+        return hidden_states
+
+
+class AllegroDownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        spatial_downsample: bool = True,
+        temporal_downsample: bool = False,
+        downsample_padding: int = 1,
+    ):
+        super().__init__()
+
+        resnets = []
+        temp_convs = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=None,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                AllegroTemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        if temporal_downsample:
+            self.temp_convs_down = AllegroTemporalConvLayer(
+                out_channels, out_channels, dropout=0.1, norm_num_groups=resnet_groups, down_sample=True, stride=3
+            )
+        self.add_temp_downsample = temporal_downsample
+
+        if spatial_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            hidden_states = resnet(hidden_states, temb=None)
+            hidden_states = temp_conv(hidden_states, batch_size=batch_size)
+
+        if self.add_temp_downsample:
+            hidden_states = self.temp_convs_down(hidden_states, batch_size=batch_size)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return hidden_states
+
+
+class AllegroUpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        spatial_upsample: bool = True,
+        temporal_upsample: bool = False,
+        temb_channels: Optional[int] = None,
+    ):
+        super().__init__()
+
+        resnets = []
+        temp_convs = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            temp_convs.append(
+                AllegroTemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        self.add_temp_upsample = temporal_upsample
+        if temporal_upsample:
+            self.temp_conv_up = AllegroTemporalConvLayer(
+                out_channels, out_channels, dropout=0.1, norm_num_groups=resnet_groups, up_sample=True, stride=3
+            )
+
+        if spatial_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            hidden_states = resnet(hidden_states, temb=None)
+            hidden_states = temp_conv(hidden_states, batch_size=batch_size)
+
+        if self.add_temp_upsample:
+            hidden_states = self.temp_conv_up(hidden_states, batch_size=batch_size)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return hidden_states
+
+
+class AllegroMidBlock3DConv(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        temp_convs = [
+            AllegroTemporalConvLayer(
+                in_channels,
+                in_channels,
+                dropout=0.1,
+                norm_num_groups=resnet_groups,
+            )
+        ]
+        attentions = []
+
+        if attention_head_dim is None:
+            attention_head_dim = in_channels
+
+        for _ in range(num_layers):
+            if add_attention:
+                attentions.append(
+                    Attention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        rescale_output_factor=output_scale_factor,
+                        eps=resnet_eps,
+                        norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None,
+                        spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
+                        residual_connection=True,
+                        bias=True,
+                        upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+            temp_convs.append(
+                AllegroTemporalConvLayer(
+                    in_channels,
+                    in_channels,
+                    dropout=0.1,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        hidden_states = self.resnets[0](hidden_states, temb=None)
+
+        hidden_states = self.temp_convs[0](hidden_states, batch_size=batch_size)
+
+        for attn, resnet, temp_conv in zip(self.attentions, self.resnets[1:], self.temp_convs[1:]):
+            hidden_states = attn(hidden_states)
+            hidden_states = resnet(hidden_states, temb=None)
+            hidden_states = temp_conv(hidden_states, batch_size=batch_size)
+
+        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return hidden_states
+
+
+class AllegroEncoder3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str, ...] = (
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        temporal_downsample_blocks: Tuple[bool, ...] = [True, True, False, False],
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+    ):
+        super().__init__()
+
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.temp_conv_in = nn.Conv3d(
+            in_channels=block_out_channels[0],
+            out_channels=block_out_channels[0],
+            kernel_size=(3, 1, 1),
+            padding=(1, 0, 0),
+        )
+
+        self.down_blocks = nn.ModuleList([])
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            if down_block_type == "AllegroDownBlock3D":
+                down_block = AllegroDownBlock3D(
+                    num_layers=layers_per_block,
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    spatial_downsample=not is_final_block,
+                    temporal_downsample=temporal_downsample_blocks[i],
+                    resnet_eps=1e-6,
+                    downsample_padding=0,
+                    resnet_act_fn=act_fn,
+                    resnet_groups=norm_num_groups,
+                )
+            else:
+                raise ValueError("Invalid `down_block_type` encountered. Must be `AllegroDownBlock3D`")
+
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = AllegroMidBlock3DConv(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default",
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=None,
+        )
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+
+        self.temp_conv_out = nn.Conv3d(block_out_channels[-1], block_out_channels[-1], (3, 1, 1), padding=(1, 0, 0))
+        self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        batch_size = sample.shape[0]
+
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_in(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        residual = sample
+        sample = self.temp_conv_in(sample)
+        sample = sample + residual
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    return module(*inputs)
+
+                return custom_forward
+
+            # Down blocks
+            for down_block in self.down_blocks:
+                sample = torch.utils.checkpoint.checkpoint(create_custom_forward(down_block), sample)
+
+            # Mid block
+            sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample)
+        else:
+            # Down blocks
+            for down_block in self.down_blocks:
+                sample = down_block(sample)
+
+            # Mid block
+            sample = self.mid_block(sample)
+
+        # Post process
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        residual = sample
+        sample = self.temp_conv_out(sample)
+        sample = sample + residual
+
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_out(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return sample
+
+
+class AllegroDecoder3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 4,
+        out_channels: int = 3,
+        up_block_types: Tuple[str, ...] = (
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+        ),
+        temporal_upsample_blocks: Tuple[bool, ...] = [False, True, True, False],
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        norm_type: str = "group",  # group, spatial
+    ):
+        super().__init__()
+
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            block_out_channels[-1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.temp_conv_in = nn.Conv3d(block_out_channels[-1], block_out_channels[-1], (3, 1, 1), padding=(1, 0, 0))
+
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        temb_channels = in_channels if norm_type == "spatial" else None
+
+        # mid
+        self.mid_block = AllegroMidBlock3DConv(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            output_scale_factor=1,
+            resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            temb_channels=temb_channels,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+
+            is_final_block = i == len(block_out_channels) - 1
+
+            if up_block_type == "AllegroUpBlock3D":
+                up_block = AllegroUpBlock3D(
+                    num_layers=layers_per_block + 1,
+                    in_channels=prev_output_channel,
+                    out_channels=output_channel,
+                    spatial_upsample=not is_final_block,
+                    temporal_upsample=temporal_upsample_blocks[i],
+                    resnet_eps=1e-6,
+                    resnet_act_fn=act_fn,
+                    resnet_groups=norm_num_groups,
+                    temb_channels=temb_channels,
+                    resnet_time_scale_shift=norm_type,
+                )
+            else:
+                raise ValueError("Invalid `UP_block_type` encountered. Must be `AllegroUpBlock3D`")
+
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_type == "spatial":
+            self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+        else:
+            self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+
+        self.conv_act = nn.SiLU()
+
+        self.temp_conv_out = nn.Conv3d(block_out_channels[0], block_out_channels[0], (3, 1, 1), padding=(1, 0, 0))
+        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        batch_size = sample.shape[0]
+
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_in(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        residual = sample
+        sample = self.temp_conv_in(sample)
+        sample = sample + residual
+
+        upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    return module(*inputs)
+
+                return custom_forward
+
+            # Mid block
+            sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample)
+
+            # Up blocks
+            for up_block in self.up_blocks:
+                sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample)
+
+        else:
+            # Mid block
+            sample = self.mid_block(sample)
+            sample = sample.to(upscale_dtype)
+
+            # Up blocks
+            for up_block in self.up_blocks:
+                sample = up_block(sample)
+
+        # Post process
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        residual = sample
+        sample = self.temp_conv_out(sample)
+        sample = sample + residual
+
+        sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        sample = self.conv_out(sample)
+
+        sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return sample
+
+
+class AutoencoderKLAllegro(ModelMixin, ConfigMixin):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used in
+    [Allegro](https://github.com/rhymes-ai/Allegro).
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        in_channels (int, defaults to `3`):
+            Number of channels in the input image.
+        out_channels (int, defaults to `3`):
+            Number of channels in the output.
+        down_block_types (`Tuple[str, ...]`, defaults to `("AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D")`):
+            Tuple of strings denoting which types of down blocks to use.
+        up_block_types (`Tuple[str, ...]`, defaults to `("AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D")`):
+            Tuple of strings denoting which types of up blocks to use.
+        block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+            Tuple of integers denoting number of output channels in each block.
+        temporal_downsample_blocks (`Tuple[bool, ...]`, defaults to `(True, True, False, False)`):
+            Tuple of booleans denoting which blocks to enable temporal downsampling in.
+        latent_channels (`int`, defaults to `4`):
+            Number of channels in latents.
+        layers_per_block (`int`, defaults to `2`):
+            Number of resnet or attention or temporal convolution layers per down/up block.
+        act_fn (`str`, defaults to `"silu"`):
+            The activation function to use.
+        norm_num_groups (`int`, defaults to `32`):
+            Number of groups to use in normalization layers.
+        temporal_compression_ratio (`int`, defaults to `4`):
+            Ratio by which temporal dimension of samples are compressed.
+        sample_size (`int`, defaults to `320`):
+            Default latent size.
+        scaling_factor (`float`, defaults to `0.13235`):
+            The component-wise standard deviation of the trained latent space computed using the first batch of the
+            training set. This is used to scale the latent space to have unit variance when training the diffusion
+            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+            Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
+        force_upcast (`bool`, default to `True`):
+            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str, ...] = (
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+            "AllegroDownBlock3D",
+        ),
+        up_block_types: Tuple[str, ...] = (
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+            "AllegroUpBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        temporal_downsample_blocks: Tuple[bool, ...] = (True, True, False, False),
+        temporal_upsample_blocks: Tuple[bool, ...] = (False, True, True, False),
+        latent_channels: int = 4,
+        layers_per_block: int = 2,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        temporal_compression_ratio: float = 4,
+        sample_size: int = 320,
+        scaling_factor: float = 0.13,
+        force_upcast: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.encoder = AllegroEncoder3D(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            temporal_downsample_blocks=temporal_downsample_blocks,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+        )
+        self.decoder = AllegroDecoder3D(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            temporal_upsample_blocks=temporal_upsample_blocks,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+        )
+        self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+        self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1)
+
+        # TODO(aryan): For the 1.0.0 refactor, `temporal_compression_ratio` can be inferred directly and we don't need
+        # to use a specific parameter here or in other VAEs.
+
+        self.use_slicing = False
+        self.use_tiling = False
+
+        self.spatial_compression_ratio = 2 ** (len(block_out_channels) - 1)
+        self.tile_overlap_t = 8
+        self.tile_overlap_h = 120
+        self.tile_overlap_w = 80
+        sample_frames = 24
+
+        self.kernel = (sample_frames, sample_size, sample_size)
+        self.stride = (
+            sample_frames - self.tile_overlap_t,
+            sample_size - self.tile_overlap_h,
+            sample_size - self.tile_overlap_w,
+        )
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (AllegroEncoder3D, AllegroDecoder3D)):
+            module.gradient_checkpointing = value
+
+    def enable_tiling(self) -> 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.
+        """
+        self.use_tiling = True
+
+    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
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        # TODO(aryan)
+        # if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
+        if self.use_tiling:
+            return self.tiled_encode(x)
+
+        raise NotImplementedError("Encoding without tiling has not been implemented yet.")
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        r"""
+        Encode a batch of videos into latents.
+
+        Args:
+            x (`torch.Tensor`):
+                Input batch of videos.
+            return_dict (`bool`, 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) -> torch.Tensor:
+        # TODO(aryan): refactor tiling implementation
+        # if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height):
+        if self.use_tiling:
+            return self.tiled_decode(z)
+
+        raise NotImplementedError("Decoding without tiling has not been implemented yet.")
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        """
+        Decode a batch of videos.
+
+        Args:
+            z (`torch.Tensor`):
+                Input batch of latent vectors.
+            return_dict (`bool`, 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) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z)
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        local_batch_size = 1
+        rs = self.spatial_compression_ratio
+        rt = self.config.temporal_compression_ratio
+
+        batch_size, num_channels, num_frames, height, width = x.shape
+
+        output_num_frames = math.floor((num_frames - self.kernel[0]) / self.stride[0]) + 1
+        output_height = math.floor((height - self.kernel[1]) / self.stride[1]) + 1
+        output_width = math.floor((width - self.kernel[2]) / self.stride[2]) + 1
+
+        count = 0
+        output_latent = x.new_zeros(
+            (
+                output_num_frames * output_height * output_width,
+                2 * self.config.latent_channels,
+                self.kernel[0] // rt,
+                self.kernel[1] // rs,
+                self.kernel[2] // rs,
+            )
+        )
+        vae_batch_input = x.new_zeros((local_batch_size, num_channels, self.kernel[0], self.kernel[1], self.kernel[2]))
+
+        for i in range(output_num_frames):
+            for j in range(output_height):
+                for k in range(output_width):
+                    n_start, n_end = i * self.stride[0], i * self.stride[0] + self.kernel[0]
+                    h_start, h_end = j * self.stride[1], j * self.stride[1] + self.kernel[1]
+                    w_start, w_end = k * self.stride[2], k * self.stride[2] + self.kernel[2]
+
+                    video_cube = x[:, :, n_start:n_end, h_start:h_end, w_start:w_end]
+                    vae_batch_input[count % local_batch_size] = video_cube
+
+                    if (
+                        count % local_batch_size == local_batch_size - 1
+                        or count == output_num_frames * output_height * output_width - 1
+                    ):
+                        latent = self.encoder(vae_batch_input)
+
+                        if (
+                            count == output_num_frames * output_height * output_width - 1
+                            and count % local_batch_size != local_batch_size - 1
+                        ):
+                            output_latent[count - count % local_batch_size :] = latent[: count % local_batch_size + 1]
+                        else:
+                            output_latent[count - local_batch_size + 1 : count + 1] = latent
+
+                        vae_batch_input = x.new_zeros(
+                            (local_batch_size, num_channels, self.kernel[0], self.kernel[1], self.kernel[2])
+                        )
+
+                    count += 1
+
+        latent = x.new_zeros(
+            (batch_size, 2 * self.config.latent_channels, num_frames // rt, height // rs, width // rs)
+        )
+        output_kernel = self.kernel[0] // rt, self.kernel[1] // rs, self.kernel[2] // rs
+        output_stride = self.stride[0] // rt, self.stride[1] // rs, self.stride[2] // rs
+        output_overlap = (
+            output_kernel[0] - output_stride[0],
+            output_kernel[1] - output_stride[1],
+            output_kernel[2] - output_stride[2],
+        )
+
+        for i in range(output_num_frames):
+            n_start, n_end = i * output_stride[0], i * output_stride[0] + output_kernel[0]
+            for j in range(output_height):
+                h_start, h_end = j * output_stride[1], j * output_stride[1] + output_kernel[1]
+                for k in range(output_width):
+                    w_start, w_end = k * output_stride[2], k * output_stride[2] + output_kernel[2]
+                    latent_mean = _prepare_for_blend(
+                        (i, output_num_frames, output_overlap[0]),
+                        (j, output_height, output_overlap[1]),
+                        (k, output_width, output_overlap[2]),
+                        output_latent[i * output_height * output_width + j * output_width + k].unsqueeze(0),
+                    )
+                    latent[:, :, n_start:n_end, h_start:h_end, w_start:w_end] += latent_mean
+
+        latent = latent.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        latent = self.quant_conv(latent)
+        latent = latent.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+        return latent
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        local_batch_size = 1
+        rs = self.spatial_compression_ratio
+        rt = self.config.temporal_compression_ratio
+
+        latent_kernel = self.kernel[0] // rt, self.kernel[1] // rs, self.kernel[2] // rs
+        latent_stride = self.stride[0] // rt, self.stride[1] // rs, self.stride[2] // rs
+
+        batch_size, num_channels, num_frames, height, width = z.shape
+
+        ## post quant conv (a mapping)
+        z = z.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        z = self.post_quant_conv(z)
+        z = z.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
+
+        output_num_frames = math.floor((num_frames - latent_kernel[0]) / latent_stride[0]) + 1
+        output_height = math.floor((height - latent_kernel[1]) / latent_stride[1]) + 1
+        output_width = math.floor((width - latent_kernel[2]) / latent_stride[2]) + 1
+
+        count = 0
+        decoded_videos = z.new_zeros(
+            (
+                output_num_frames * output_height * output_width,
+                self.config.out_channels,
+                self.kernel[0],
+                self.kernel[1],
+                self.kernel[2],
+            )
+        )
+        vae_batch_input = z.new_zeros(
+            (local_batch_size, num_channels, latent_kernel[0], latent_kernel[1], latent_kernel[2])
+        )
+
+        for i in range(output_num_frames):
+            for j in range(output_height):
+                for k in range(output_width):
+                    n_start, n_end = i * latent_stride[0], i * latent_stride[0] + latent_kernel[0]
+                    h_start, h_end = j * latent_stride[1], j * latent_stride[1] + latent_kernel[1]
+                    w_start, w_end = k * latent_stride[2], k * latent_stride[2] + latent_kernel[2]
+
+                    current_latent = z[:, :, n_start:n_end, h_start:h_end, w_start:w_end]
+                    vae_batch_input[count % local_batch_size] = current_latent
+
+                    if (
+                        count % local_batch_size == local_batch_size - 1
+                        or count == output_num_frames * output_height * output_width - 1
+                    ):
+                        current_video = self.decoder(vae_batch_input)
+
+                        if (
+                            count == output_num_frames * output_height * output_width - 1
+                            and count % local_batch_size != local_batch_size - 1
+                        ):
+                            decoded_videos[count - count % local_batch_size :] = current_video[
+                                : count % local_batch_size + 1
+                            ]
+                        else:
+                            decoded_videos[count - local_batch_size + 1 : count + 1] = current_video
+
+                        vae_batch_input = z.new_zeros(
+                            (local_batch_size, num_channels, latent_kernel[0], latent_kernel[1], latent_kernel[2])
+                        )
+
+                    count += 1
+
+        video = z.new_zeros((batch_size, self.config.out_channels, num_frames * rt, height * rs, width * rs))
+        video_overlap = (
+            self.kernel[0] - self.stride[0],
+            self.kernel[1] - self.stride[1],
+            self.kernel[2] - self.stride[2],
+        )
+
+        for i in range(output_num_frames):
+            n_start, n_end = i * self.stride[0], i * self.stride[0] + self.kernel[0]
+            for j in range(output_height):
+                h_start, h_end = j * self.stride[1], j * self.stride[1] + self.kernel[1]
+                for k in range(output_width):
+                    w_start, w_end = k * self.stride[2], k * self.stride[2] + self.kernel[2]
+                    out_video_blend = _prepare_for_blend(
+                        (i, output_num_frames, video_overlap[0]),
+                        (j, output_height, video_overlap[1]),
+                        (k, output_width, video_overlap[2]),
+                        decoded_videos[i * output_height * output_width + j * output_width + k].unsqueeze(0),
+                    )
+                    video[:, :, n_start:n_end, h_start:h_end, w_start:w_end] += out_video_blend
+
+        video = video.permute(0, 2, 1, 3, 4).contiguous()
+        return video
+
+    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.
+            generator (`torch.Generator`, *optional*):
+                PyTorch random number generator.
+        """
+        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)
+
+
+def _prepare_for_blend(n_param, h_param, w_param, x):
+    # TODO(aryan): refactor
+    n, n_max, overlap_n = n_param
+    h, h_max, overlap_h = h_param
+    w, w_max, overlap_w = w_param
+    if overlap_n > 0:
+        if n > 0:  # the head overlap part decays from 0 to 1
+            x[:, :, 0:overlap_n, :, :] = x[:, :, 0:overlap_n, :, :] * (
+                torch.arange(0, overlap_n).float().to(x.device) / overlap_n
+            ).reshape(overlap_n, 1, 1)
+        if n < n_max - 1:  # the tail overlap part decays from 1 to 0
+            x[:, :, -overlap_n:, :, :] = x[:, :, -overlap_n:, :, :] * (
+                1 - torch.arange(0, overlap_n).float().to(x.device) / overlap_n
+            ).reshape(overlap_n, 1, 1)
+    if h > 0:
+        x[:, :, :, 0:overlap_h, :] = x[:, :, :, 0:overlap_h, :] * (
+            torch.arange(0, overlap_h).float().to(x.device) / overlap_h
+        ).reshape(overlap_h, 1)
+    if h < h_max - 1:
+        x[:, :, :, -overlap_h:, :] = x[:, :, :, -overlap_h:, :] * (
+            1 - torch.arange(0, overlap_h).float().to(x.device) / overlap_h
+        ).reshape(overlap_h, 1)
+    if w > 0:
+        x[:, :, :, :, 0:overlap_w] = x[:, :, :, :, 0:overlap_w] * (
+            torch.arange(0, overlap_w).float().to(x.device) / overlap_w
+        )
+    if w < w_max - 1:
+        x[:, :, :, :, -overlap_w:] = x[:, :, :, :, -overlap_w:] * (
+            1 - torch.arange(0, overlap_w).float().to(x.device) / overlap_w
+        )
+    return x