src/models/transformers/wan22.py

"""
Modified Wan2.2 Diffusion Transformer for LoomVideo.

Extends the HuggingFace Diffusers Wan transformer to support:
- Separate self-attention and cross-attention execution paths
- Variable-length sequence handling with cu_seq_lens
- Source video conditioning via learnable patch embedding
- Reference image/video conditioning with custom temporal RoPE offsets

Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py
"""

import math
from typing import Optional, Tuple, List

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from diffusers.configuration_utils import register_to_config
from diffusers.utils import logging
from diffusers.models.attention import FeedForward
from diffusers.models.embeddings import get_1d_rotary_pos_embed
from diffusers.models.normalization import FP32LayerNorm
from diffusers.models.transformers.transformer_wan import (
    WanAttention,
    WanRotaryPosEmbed,
    WanTimeTextImageEmbedding,
    WanAttnProcessor,
    WanTransformerBlock,
    WanTransformer3DModel,
    _get_qkv_projections,
)

logger = logging.get_logger(__name__)


class WanAttnProcessor(WanAttnProcessor):
    """Attention processor with QKV extraction for cross-attention in LoomVideo."""

    _attention_backend = None

    def __init__(self):
        super().__init__()
        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError("WanAttnProcessor requires PyTorch 2.0+.")

    def get_qkv(
        self,
        attn: "WanAttention",
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
        """
        Extract Q, K, V projections with optional rotary embeddings.

        Returns:
            Tuple of (query, key, value, encoder_hidden_states, encoder_hidden_states_img).
        """
        encoder_hidden_states_img = None
        if attn.add_k_proj is not None:
            # 512 is the context length of the text encoder
            image_context_length = encoder_hidden_states.shape[1] - 512
            encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length]
            encoder_hidden_states = encoder_hidden_states[:, image_context_length:]

        query, key, value = _get_qkv_projections(
            attn, hidden_states, encoder_hidden_states
        )

        query = attn.norm_q(query)
        key = attn.norm_k(key)

        query = query.unflatten(2, (attn.heads, -1))
        key = key.unflatten(2, (attn.heads, -1))
        value = value.unflatten(2, (attn.heads, -1))

        if rotary_emb is not None:

            def apply_rotary_emb(
                hidden_states: torch.Tensor,
                freqs_cos: torch.Tensor,
                freqs_sin: torch.Tensor,
            ):
                x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
                cos = freqs_cos[..., 0::2]
                sin = freqs_sin[..., 1::2]
                out = torch.empty_like(hidden_states)
                out[..., 0::2] = x1 * cos - x2 * sin
                out[..., 1::2] = x1 * sin + x2 * cos
                return out.type_as(hidden_states)

            query = apply_rotary_emb(query, *rotary_emb)
            key = apply_rotary_emb(key, *rotary_emb)

        # [B, L, Num_heads, Head_dims] -> [B, Num_heads, L, Head_dims]
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)

        return query, key, value, encoder_hidden_states, encoder_hidden_states_img


class WanTransformerBlock(WanTransformerBlock):
    """
    Extended transformer block with separate self-attention and cross-attention paths.

    Adds methods for split execution:
    - forward_selfattn: self-attention only (for training with flash cross-attn)
    - forward_crossattn_later_layer: cross-attn + FFN (paired with forward_selfattn)
    """

    def __init__(
        self,
        dim: int,
        ffn_dim: int,
        num_heads: int,
        qk_norm: str = "rms_norm_across_heads",
        cross_attn_norm: bool = False,
        eps: float = 1e-6,
        added_kv_proj_dim: Optional[int] = None,
    ):
        super().__init__(
            dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim
        )

        self.dim = dim
        self.num_heads = num_heads
        self.eps = eps

        # 1. Self-attention
        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
        self.attn1 = WanAttention(
            dim=dim,
            heads=num_heads,
            dim_head=dim // num_heads,
            eps=eps,
            cross_attention_dim_head=None,
            processor=WanAttnProcessor(),
        )

        # 2. Cross-attention
        self.attn2 = WanAttention(
            dim=dim,
            heads=num_heads,
            dim_head=dim // num_heads,
            eps=eps,
            added_kv_proj_dim=added_kv_proj_dim,
            cross_attention_dim_head=dim // num_heads,
            processor=WanAttnProcessor(),
        )
        self.norm2 = (
            FP32LayerNorm(dim, eps, elementwise_affine=True)
            if cross_attn_norm
            else nn.Identity()
        )

        # 3. Feed-forward
        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)

    def _parse_temb(self, temb: torch.Tensor):
        """Parse timestep embedding into shift/scale/gate components."""
        if temb.ndim == 4:
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
                self.scale_shift_table.unsqueeze(0) + temb.float()
            ).chunk(6, dim=2)
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
            gate_msa = gate_msa.squeeze(2)
            c_shift_msa = c_shift_msa.squeeze(2)
            c_scale_msa = c_scale_msa.squeeze(2)
            c_gate_msa = c_gate_msa.squeeze(2)
        else:
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
                self.scale_shift_table + temb.float()
            ).chunk(6, dim=1)
        return shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa

    def forward_selfattn(
        self,
        hidden_states: torch.Tensor,
        temb: torch.Tensor,
        rotary_emb: torch.Tensor,
        cu_seq_lens: list = None,
    ) -> torch.Tensor:
        """
        Execute only the self-attention portion of this block.

        Args:
            hidden_states: Input tensor.
            temb: Timestep embedding (adaptive normalization parameters).
            rotary_emb: Rotary position embeddings.
            cu_seq_lens: Cumulative sequence lengths for per-sample attention.

        Returns:
            Hidden states after self-attention (before cross-attention and FFN).
        """
        shift_msa, scale_msa, gate_msa, _, _, _ = self._parse_temb(temb)

        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)

        if cu_seq_lens is not None and len(cu_seq_lens) > 2:
            # Per-sample self-attention for variable-length sequences
            attn_outputs = []
            for i in range(len(cu_seq_lens) - 1):
                start = int(cu_seq_lens[i])
                end = int(cu_seq_lens[i + 1])
                sample_norm_hidden = norm_hidden_states[:, start:end, :]

                if isinstance(rotary_emb, (list, tuple)):
                    sample_rotary = [r[:, start:end, ...] if r is not None else None for r in rotary_emb]
                elif rotary_emb is not None:
                    sample_rotary = rotary_emb[:, start:end, ...]
                else:
                    sample_rotary = None
                sample_attn_output = self.attn1(sample_norm_hidden, None, None, sample_rotary)
                attn_outputs.append(sample_attn_output)

            attn_output = torch.cat(attn_outputs, dim=1)
        else:
            attn_output = self.attn1(norm_hidden_states, None, None, rotary_emb)

        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)

        return hidden_states

    def forward_crossattn_later_layer(
        self,
        hidden_states: torch.Tensor,
        attn_output: torch.Tensor,
        temb: torch.Tensor,
    ) -> torch.Tensor:
        """
        Execute cross-attention residual addition and feed-forward network.

        Args:
            hidden_states: Current hidden states (after self-attention).
            attn_output: Cross-attention output to be added as residual.
            temb: Timestep embedding for adaptive normalization.

        Returns:
            Hidden states after cross-attention residual and FFN.
        """
        _, _, _, c_shift_msa, c_scale_msa, c_gate_msa = self._parse_temb(temb)

        hidden_states = hidden_states + attn_output

        # Feed-forward network
        norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states)
        ff_output = self.ffn(norm_hidden_states)
        hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)

        return hidden_states


class WanTransformer3DModel(WanTransformer3DModel):
    """
    Extended Wan 3D Transformer with source/reference conditioning support.

    Adds:
    - Source video conditioning via learnable patch embedding
    - Reference conditioning with custom temporal RoPE offsets
    - Separate early-layer and output-projection methods for flexible fusion
    """

    _supports_gradient_checkpointing = True
    _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"]
    _no_split_modules = ["WanTransformerBlock"]
    _keep_in_fp32_modules = [
        "time_embedder",
        "scale_shift_table",
        "norm1",
        "norm2",
        "norm3",
    ]
    _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
    _repeated_blocks = ["WanTransformerBlock"]

    @register_to_config
    def __init__(
        self,
        patch_size: Tuple[int] = (1, 2, 2),
        num_attention_heads: int = 40,
        attention_head_dim: int = 128,
        in_channels: int = 16,
        out_channels: int = 16,
        text_dim: int = 4096,
        freq_dim: int = 256,
        ffn_dim: int = 13824,
        num_layers: int = 40,
        cross_attn_norm: bool = True,
        qk_norm: Optional[str] = "rms_norm_across_heads",
        eps: float = 1e-6,
        image_dim: Optional[int] = None,
        added_kv_proj_dim: Optional[int] = None,
        rope_max_seq_len: int = 1024,
        pos_embed_seq_len: Optional[int] = None,
    ):
        super().__init__(
            patch_size,
            num_attention_heads,
            attention_head_dim,
            in_channels,
            out_channels,
            text_dim,
            freq_dim,
            ffn_dim,
            num_layers,
            cross_attn_norm,
            qk_norm,
            eps,
            image_dim,
            added_kv_proj_dim,
            rope_max_seq_len,
            pos_embed_seq_len,
        )

        inner_dim = num_attention_heads * attention_head_dim
        out_channels = out_channels or in_channels
        self.in_channels = in_channels
        self.ffn_dim = ffn_dim
        self.qk_norm = qk_norm
        self.cross_attn_norm = cross_attn_norm
        self.attention_head_dim = attention_head_dim
        self.num_attention_heads = num_attention_heads
        self.eps = eps
        self.added_kv_proj_dim = added_kv_proj_dim
        self.num_layers = num_layers
        self.inner_dim = inner_dim

        # 1. Patch & position embedding
        self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
        self.patch_embedding = nn.Conv3d(
            in_channels, inner_dim, kernel_size=patch_size, stride=patch_size
        )

        # 2. Condition embeddings
        self.condition_embedder = WanTimeTextImageEmbedding(
            dim=inner_dim,
            time_freq_dim=freq_dim,
            time_proj_dim=inner_dim * 6,
            text_embed_dim=text_dim,
            image_embed_dim=image_dim,
            pos_embed_seq_len=pos_embed_seq_len,
        )

        # 3. Transformer blocks
        self.blocks = nn.ModuleList(
            [
                WanTransformerBlock(
                    inner_dim,
                    ffn_dim,
                    num_attention_heads,
                    qk_norm,
                    cross_attn_norm,
                    eps,
                    added_kv_proj_dim,
                )
                for _ in range(num_layers)
            ]
        )

        # 4. Output norm & projection
        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
        self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
        self.scale_shift_table = nn.Parameter(
            torch.randn(1, 2, inner_dim) / inner_dim**0.5
        )

        self.gradient_checkpointing = True

    def compute_ref_rotary_emb(
        self,
        ref_shape: Tuple[int, ...],
        time_offset: int,
        device: torch.device,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute 3D RoPE for a reference latent with a custom temporal offset.

        Reference frames share the same temporal position (given by time_offset),
        while spatial dimensions use normal indices starting from 0.

        Args:
            ref_shape: (B, C, T, H, W) shape of the reference latent tensor.
            time_offset: Temporal position index for all frames (e.g., -10, -20).
            device: Target device.

        Returns:
            Tuple of (freqs_cos, freqs_sin) tensors.
        """
        _, _, num_frames, height, width = ref_shape
        p_t, p_h, p_w = self.config.patch_size
        ppf = num_frames // p_t
        pph = height // p_h
        ppw = width // p_w

        rope_module = self.rope
        freqs_dtype = torch.float64

        # Temporal: all frames share the same time_offset
        time_pos = np.array([time_offset] * ppf, dtype=np.float64)
        freq_cos_t, freq_sin_t = get_1d_rotary_pos_embed(
            rope_module.t_dim, time_pos, theta=10000.0,
            use_real=True, repeat_interleave_real=True, freqs_dtype=freqs_dtype,
        )

        # Spatial: normal indices [0, 1, ...]
        split_sizes = [rope_module.t_dim, rope_module.h_dim, rope_module.w_dim]
        precomputed_cos = rope_module.freqs_cos.split(split_sizes, dim=1)
        precomputed_sin = rope_module.freqs_sin.split(split_sizes, dim=1)

        freq_cos_h = precomputed_cos[1][:pph]
        freq_sin_h = precomputed_sin[1][:pph]
        freq_cos_w = precomputed_cos[2][:ppw]
        freq_sin_w = precomputed_sin[2][:ppw]

        # Broadcast to (ppf, pph, ppw, dim_*)
        freq_cos_t = freq_cos_t.to(device).view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
        freq_cos_h = freq_cos_h.to(device).view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
        freq_cos_w = freq_cos_w.to(device).view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)

        freq_sin_t = freq_sin_t.to(device).view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1)
        freq_sin_h = freq_sin_h.to(device).view(1, pph, 1, -1).expand(ppf, pph, ppw, -1)
        freq_sin_w = freq_sin_w.to(device).view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1)

        freqs_cos = torch.cat([freq_cos_t, freq_cos_h, freq_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
        freqs_sin = torch.cat([freq_sin_t, freq_sin_h, freq_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)

        return freqs_cos, freqs_sin

    def forward_early_layers(
        self,
        hidden_states: List[torch.Tensor],
        timestep: torch.LongTensor,
        encoder_hidden_states: torch.Tensor,
        source_hidden_states: Optional[List[torch.Tensor]] = None,
        source_scale: Optional[torch.Tensor] = None,
        ref_hidden_states: Optional[List[torch.Tensor]] = None,
    ):
        """
        Process patch embedding, timestep conditioning, source/ref concatenation.

        Prepares all inputs for the main transformer blocks by:
        1. Applying patch embedding to each latent sample
        2. Adding source conditioning (scaled by timestep)
        3. Concatenating reference latents with custom temporal RoPE
        4. Computing timestep projections

        Args:
            hidden_states: List of latent tensors (None for non-generation samples).
            timestep: Diffusion timestep for each generation sample.
            encoder_hidden_states: T5 text encoder hidden states.
            source_hidden_states: Optional source video latents.
            source_scale: Timestep-dependent scale for source conditioning.
            ref_hidden_states: Optional list of reference latent lists.

        Returns:
            Tuple of processed tensors and metadata for the transformer blocks.
        """
        hidden_states_list = []
        shape_list = []
        rotary_emb_cos = []
        rotary_emb_sin = []
        cu_seq_lens = [0]
        sample_index = []
        seq_lens = []
        gen_seq_lens = []

        for index, hidden_state in enumerate(hidden_states):
            if hidden_state is not None:
                valid_sample_idx = len(seq_lens)
                shape = hidden_state.shape

                hidden_state = self.patch_embedding(hidden_state)
                hidden_state = hidden_state.flatten(2).transpose(1, 2)  # [1, L, C]

                # Add source video conditioning
                if (source_hidden_states is not None
                        and source_hidden_states[index] is not None
                        and source_scale is not None
                        and hasattr(self, 'source_patch_embedding')):
                    source_input = source_hidden_states[index]
                    source_encoded = self.source_patch_embedding(source_input)
                    source_encoded = source_encoded.flatten(2).transpose(1, 2)
                    source_scale_idx = index if source_scale.shape[0] == len(hidden_states) else valid_sample_idx
                    hidden_state = hidden_state + source_encoded * source_scale[source_scale_idx]

                # Record original sequence length (for loss computation, excluding ref tokens)
                original_seq_len = hidden_state.shape[1]
                gen_seq_lens.append(original_seq_len)

                # Compute rotary embeddings for the main latent
                fake_tensor = torch.empty(shape, device=hidden_state.device, dtype=hidden_state.dtype)
                rotary_emb = self.rope(fake_tensor)

                # Concatenate reference latents with negative temporal offsets
                if (ref_hidden_states is not None and ref_hidden_states[index] is not None):
                    ref_list = ref_hidden_states[index]
                    for ref_idx, ref_input in enumerate(ref_list):
                        ref_encoded = self.patch_embedding(ref_input)
                        ref_encoded = ref_encoded.flatten(2).transpose(1, 2)
                        hidden_state = torch.cat([hidden_state, ref_encoded], dim=1)

                        # Negative temporal offset: -10, -20, -30, ... for each ref
                        ref_time_offset = -10 * (ref_idx + 1)
                        ref_rotary_emb = self.compute_ref_rotary_emb(
                            ref_shape=ref_input.shape,
                            time_offset=ref_time_offset,
                            device=hidden_state.device,
                        )
                        rotary_emb = (
                            torch.cat([rotary_emb[0], ref_rotary_emb[0]], dim=1),
                            torch.cat([rotary_emb[1], ref_rotary_emb[1]], dim=1),
                        )

                hidden_states_list.append(hidden_state)
                rotary_emb_cos.append(rotary_emb[0])
                rotary_emb_sin.append(rotary_emb[1])
                cu_seq_lens.append(cu_seq_lens[-1] + hidden_state.shape[1])
                sample_index.append(index)
                shape_list.append(shape)
                seq_lens.append(hidden_state.shape[1])

        hidden_states = torch.cat(hidden_states_list, dim=1)  # [1, sum(L), C]
        rotary_emb_cos = torch.cat(rotary_emb_cos, dim=1)
        rotary_emb_sin = torch.cat(rotary_emb_sin, dim=1)
        rotary_emb = (rotary_emb_cos, rotary_emb_sin)

        # Timestep conditioning
        if timestep.ndim == 2:
            ts_seq_len = timestep.shape[1]
            timestep = timestep.flatten()
        else:
            ts_seq_len = None

        temb, timestep_proj, encoder_hidden_states, _ = (
            self.condition_embedder(
                timestep,
                encoder_hidden_states,
                timestep_seq_len=ts_seq_len,
            )
        )
        if ts_seq_len is not None:
            timestep_proj = timestep_proj.unflatten(2, (6, -1))
        else:
            timestep_proj = timestep_proj.unflatten(1, (6, -1))

        # Expand timestep projections to match sequence lengths (zero for ref tokens)
        num_valid_samples = len(seq_lens)
        if timestep_proj.shape[0] == num_valid_samples and num_valid_samples > 0:
            expanded_timestep_proj = []
            expanded_temb = []
            for i, length in enumerate(seq_lens):
                gen_len = gen_seq_lens[i]
                ref_len = length - gen_len
                if ts_seq_len is None:
                    expanded_timestep_proj.append(timestep_proj[i:i + 1].expand(gen_len, -1, -1))
                    if temb is not None:
                        expanded_temb.append(temb[i:i + 1].expand(gen_len, -1))
                    # Zero timestep embedding for ref tokens
                    if ref_len > 0:
                        expanded_timestep_proj.append(
                            torch.zeros(ref_len, timestep_proj.shape[1], timestep_proj.shape[2],
                                        device=timestep_proj.device, dtype=timestep_proj.dtype)
                        )
                        if temb is not None:
                            expanded_temb.append(
                                torch.zeros(ref_len, temb.shape[1], device=temb.device, dtype=temb.dtype)
                            )
                else:
                    expanded_timestep_proj.append(timestep_proj[i])
                    if temb is not None:
                        expanded_temb.append(temb[i])

            if ts_seq_len is None:
                timestep_proj = torch.cat(expanded_timestep_proj, dim=0).unsqueeze(0)
                if temb is not None:
                    temb = torch.cat(expanded_temb, dim=0).unsqueeze(0)
            else:
                timestep_proj = torch.cat(expanded_timestep_proj, dim=0).unsqueeze(0)
                if temb is not None:
                    temb = torch.cat(expanded_temb, dim=0).unsqueeze(0)

        return (
            hidden_states,
            encoder_hidden_states,
            timestep_proj,
            rotary_emb,
            temb,
            shape_list,
            cu_seq_lens,
            sample_index,
            gen_seq_lens,
        )

    def get_output(self, hidden_states: torch.Tensor, temb: torch.Tensor, cu_seq_lens, shape, gen_seq_lens=None):
        """
        Project hidden states back to pixel space via unpatchify.

        Args:
            hidden_states: Transformer output tensor.
            temb: Timestep embedding for final adaptive normalization.
            cu_seq_lens: Cumulative sequence lengths.
            shape: Original latent shapes per sample.
            gen_seq_lens: Original sequence lengths (excluding ref tokens).
                If provided, only these tokens are projected (ref tokens discarded).

        Returns:
            List of output tensors, one per sample.
        """
        output_list = []

        for index in range(len(cu_seq_lens) - 1):
            start_idx = cu_seq_lens[index]
            # Only project original tokens (exclude ref tokens)
            if gen_seq_lens is not None:
                end_idx = cu_seq_lens[index] + gen_seq_lens[index]
            else:
                end_idx = cu_seq_lens[index + 1]
            hidden_state = hidden_states[:, start_idx:end_idx, :]

            if temb.ndim == 3 and temb.shape[1] == hidden_states.shape[1]:
                sample_temb = temb[:, start_idx:end_idx, :]
            elif temb.ndim == 2 and temb.shape[0] > 1:
                sample_temb = temb[index:index + 1]
            else:
                sample_temb = temb

            batch_size, num_channels, num_frames, height, width = shape[index]
            p_t, p_h, p_w = self.config.patch_size
            post_patch_num_frames = num_frames // p_t
            post_patch_height = height // p_h
            post_patch_width = width // p_w

            # Adaptive normalization for output
            if sample_temb.ndim == 3:
                shift, scale = (
                    self.scale_shift_table.unsqueeze(0) + sample_temb.unsqueeze(2)
                ).chunk(2, dim=2)
                shift = shift.squeeze(2)
                scale = scale.squeeze(2)
            else:
                shift, scale = (self.scale_shift_table + sample_temb.unsqueeze(1)).chunk(2, dim=1)

            shift = shift.to(hidden_state.device)
            scale = scale.to(hidden_state.device)

            hidden_state = (
                self.norm_out(hidden_state.float()) * (1 + scale) + shift
            ).type_as(hidden_state)
            hidden_state = self.proj_out(hidden_state)

            # Unpatchify: reshape back to video dimensions
            hidden_state = hidden_state.reshape(
                batch_size,
                post_patch_num_frames,
                post_patch_height,
                post_patch_width,
                p_t,
                p_h,
                p_w,
                -1,
            )
            hidden_state = hidden_state.permute(0, 7, 1, 4, 2, 5, 3, 6)
            output = hidden_state.flatten(6, 7).flatten(4, 5).flatten(2, 3)
            output_list.append(output)

        return output_list