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