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)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py