class FreeNoiseTransformerBlock(nn.Module):
    r"""
    A FreeNoise Transformer block.

    Parameters:
        dim (`int`):
            The number of channels in the input and output.
        num_attention_heads (`int`):
            The number of heads to use for multi-head attention.
        attention_head_dim (`int`):
            The number of channels in each head.
        dropout (`float`, *optional*, defaults to 0.0):
            The dropout probability to use.
        cross_attention_dim (`int`, *optional*):
            The size of the encoder_hidden_states vector for cross attention.
        activation_fn (`str`, *optional*, defaults to `"geglu"`):
            Activation function to be used in feed-forward.
        num_embeds_ada_norm (`int`, *optional*):
            The number of diffusion steps used during training. See `Transformer2DModel`.
        attention_bias (`bool`, defaults to `False`):
            Configure if the attentions should contain a bias parameter.
        only_cross_attention (`bool`, defaults to `False`):
            Whether to use only cross-attention layers. In this case two cross attention layers are used.
        double_self_attention (`bool`, defaults to `False`):
            Whether to use two self-attention layers. In this case no cross attention layers are used.
        upcast_attention (`bool`, defaults to `False`):
            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
        norm_elementwise_affine (`bool`, defaults to `True`):
            Whether to use learnable elementwise affine parameters for normalization.
        norm_type (`str`, defaults to `"layer_norm"`):
            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
        final_dropout (`bool` defaults to `False`):
            Whether to apply a final dropout after the last feed-forward layer.
        attention_type (`str`, defaults to `"default"`):
            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
        positional_embeddings (`str`, *optional*):
            The type of positional embeddings to apply to.
        num_positional_embeddings (`int`, *optional*, defaults to `None`):
            The maximum number of positional embeddings to apply.
        ff_inner_dim (`int`, *optional*):
            Hidden dimension of feed-forward MLP.
        ff_bias (`bool`, defaults to `True`):
            Whether or not to use bias in feed-forward MLP.
        attention_out_bias (`bool`, defaults to `True`):
            Whether or not to use bias in attention output project layer.
        context_length (`int`, defaults to `16`):
            The maximum number of frames that the FreeNoise block processes at once.
        context_stride (`int`, defaults to `4`):
            The number of frames to be skipped before starting to process a new batch of `context_length` frames.
        weighting_scheme (`str`, defaults to `"pyramid"`):
            The weighting scheme to use for weighting averaging of processed latent frames. As described in the
            Equation 9. of the [FreeNoise](https://arxiv.org/abs/2310.15169) paper, "pyramid" is the default setting
            used.
    """

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        dropout: float = 0.0,
        cross_attention_dim: Optional[int] = None,
        activation_fn: str = "geglu",
        num_embeds_ada_norm: Optional[int] = None,
        attention_bias: bool = False,
        only_cross_attention: bool = False,
        double_self_attention: bool = False,
        upcast_attention: bool = False,
        norm_elementwise_affine: bool = True,
        norm_type: str = "layer_norm",
        norm_eps: float = 1e-5,
        final_dropout: bool = False,
        positional_embeddings: Optional[str] = None,
        num_positional_embeddings: Optional[int] = None,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
        context_length: int = 16,
        context_stride: int = 4,
        weighting_scheme: str = "pyramid",
    ):
        super().__init__()
        self.dim = dim
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        self.dropout = dropout
        self.cross_attention_dim = cross_attention_dim
        self.activation_fn = activation_fn
        self.attention_bias = attention_bias
        self.double_self_attention = double_self_attention
        self.norm_elementwise_affine = norm_elementwise_affine
        self.positional_embeddings = positional_embeddings
        self.num_positional_embeddings = num_positional_embeddings
        self.only_cross_attention = only_cross_attention

        self.set_free_noise_properties(context_length, context_stride, weighting_scheme)

        # We keep these boolean flags for backward-compatibility.
        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
        self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
        self.use_layer_norm = norm_type == "layer_norm"
        self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"

        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
            raise ValueError(
                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
            )

        self.norm_type = norm_type
        self.num_embeds_ada_norm = num_embeds_ada_norm

        if positional_embeddings and (num_positional_embeddings is None):
            raise ValueError(
                "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
            )

        if positional_embeddings == "sinusoidal":
            self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
        else:
            self.pos_embed = None

        # Define 3 blocks. Each block has its own normalization layer.
        # 1. Self-Attn
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)

        self.attn1 = Attention(
            query_dim=dim,
            heads=num_attention_heads,
            dim_head=attention_head_dim,
            dropout=dropout,
            bias=attention_bias,
            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
            upcast_attention=upcast_attention,
            out_bias=attention_out_bias,
        )

        # 2. Cross-Attn
        if cross_attention_dim is not None or double_self_attention:
            self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)

            self.attn2 = Attention(
                query_dim=dim,
                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
                heads=num_attention_heads,
                dim_head=attention_head_dim,
                dropout=dropout,
                bias=attention_bias,
                upcast_attention=upcast_attention,
                out_bias=attention_out_bias,
            )  # is self-attn if encoder_hidden_states is none

        # 3. Feed-forward
        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )

        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)

        # let chunk size default to None
        self._chunk_size = None
        self._chunk_dim = 0

    def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]:
        frame_indices = []
        for i in range(0, num_frames - self.context_length + 1, self.context_stride):
            window_start = i
            window_end = min(num_frames, i + self.context_length)
            frame_indices.append((window_start, window_end))
        return frame_indices

    def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> List[float]:
        if weighting_scheme == "flat":
            weights = [1.0] * num_frames

        elif weighting_scheme == "pyramid":
            if num_frames % 2 == 0:
                # num_frames = 4 => [1, 2, 2, 1]
                mid = num_frames // 2
                weights = list(range(1, mid + 1))
                weights = weights + weights[::-1]
            else:
                # num_frames = 5 => [1, 2, 3, 2, 1]
                mid = (num_frames + 1) // 2
                weights = list(range(1, mid))
                weights = weights + [mid] + weights[::-1]

        elif weighting_scheme == "delayed_reverse_sawtooth":
            if num_frames % 2 == 0:
                # num_frames = 4 => [0.01, 2, 2, 1]
                mid = num_frames // 2
                weights = [0.01] * (mid - 1) + [mid]
                weights = weights + list(range(mid, 0, -1))
            else:
                # num_frames = 5 => [0.01, 0.01, 3, 2, 1]
                mid = (num_frames + 1) // 2
                weights = [0.01] * mid
                weights = weights + list(range(mid, 0, -1))
        else:
            raise ValueError(f"Unsupported value for weighting_scheme={weighting_scheme}")

        return weights

    def set_free_noise_properties(
        self, context_length: int, context_stride: int, weighting_scheme: str = "pyramid"
    ) -> None:
        self.context_length = context_length
        self.context_stride = context_stride
        self.weighting_scheme = weighting_scheme

    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> None:
        # Sets chunk feed-forward
        self._chunk_size = chunk_size
        self._chunk_dim = dim

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        cross_attention_kwargs: Dict[str, Any] = None,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        if cross_attention_kwargs is not None:
            if cross_attention_kwargs.get("scale", None) is not None:
                logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")

        cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}

        # hidden_states: [B x H x W, F, C]
        device = hidden_states.device
        dtype = hidden_states.dtype

        num_frames = hidden_states.size(1)
        frame_indices = self._get_frame_indices(num_frames)
        frame_weights = self._get_frame_weights(self.context_length, self.weighting_scheme)
        frame_weights = torch.tensor(frame_weights, device=device, dtype=dtype).unsqueeze(0).unsqueeze(-1)
        is_last_frame_batch_complete = frame_indices[-1][1] == num_frames

        # Handle out-of-bounds case if num_frames isn't perfectly divisible by context_length
        # For example, num_frames=25, context_length=16, context_stride=4, then we expect the ranges:
        #    [(0, 16), (4, 20), (8, 24), (10, 26)]
        if not is_last_frame_batch_complete:
            if num_frames < self.context_length:
                raise ValueError(f"Expected {num_frames=} to be greater or equal than {self.context_length=}")
            last_frame_batch_length = num_frames - frame_indices[-1][1]
            frame_indices.append((num_frames - self.context_length, num_frames))

        num_times_accumulated = torch.zeros((1, num_frames, 1), device=device)
        accumulated_values = torch.zeros_like(hidden_states)

        for i, (frame_start, frame_end) in enumerate(frame_indices):
            # The reason for slicing here is to ensure that if (frame_end - frame_start) is to handle
            # cases like frame_indices=[(0, 16), (16, 20)], if the user provided a video with 19 frames, or
            # essentially a non-multiple of `context_length`.
            weights = torch.ones_like(num_times_accumulated[:, frame_start:frame_end])
            weights *= frame_weights

            hidden_states_chunk = hidden_states[:, frame_start:frame_end]

            # Notice that normalization is always applied before the real computation in the following blocks.
            # 1. Self-Attention
            norm_hidden_states = self.norm1(hidden_states_chunk)

            if self.pos_embed is not None:
                norm_hidden_states = self.pos_embed(norm_hidden_states)

            attn_output = self.attn1(
                norm_hidden_states,
                encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
                attention_mask=attention_mask,
                **cross_attention_kwargs,
            )

            hidden_states_chunk = attn_output + hidden_states_chunk
            if hidden_states_chunk.ndim == 4:
                hidden_states_chunk = hidden_states_chunk.squeeze(1)

            # 2. Cross-Attention
            if self.attn2 is not None:
                norm_hidden_states = self.norm2(hidden_states_chunk)

                if self.pos_embed is not None and self.norm_type != "ada_norm_single":
                    norm_hidden_states = self.pos_embed(norm_hidden_states)

                attn_output = self.attn2(
                    norm_hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=encoder_attention_mask,
                    **cross_attention_kwargs,
                )
                hidden_states_chunk = attn_output + hidden_states_chunk

            if i == len(frame_indices) - 1 and not is_last_frame_batch_complete:
                accumulated_values[:, -last_frame_batch_length:] += (
                    hidden_states_chunk[:, -last_frame_batch_length:] * weights[:, -last_frame_batch_length:]
                )
                num_times_accumulated[:, -last_frame_batch_length:] += weights[:, -last_frame_batch_length]
            else:
                accumulated_values[:, frame_start:frame_end] += hidden_states_chunk * weights
                num_times_accumulated[:, frame_start:frame_end] += weights

        # TODO(aryan): Maybe this could be done in a better way.
        #
        # Previously, this was:
        # hidden_states = torch.where(
        #    num_times_accumulated > 0, accumulated_values / num_times_accumulated, accumulated_values
        # )
        #
        # The reasoning for the change here is `torch.where` became a bottleneck at some point when golfing memory
        # spikes. It is particularly noticeable when the number of frames is high. My understanding is that this comes
        # from tensors being copied - which is why we resort to spliting and concatenating here. I've not particularly
        # looked into this deeply because other memory optimizations led to more pronounced reductions.
        hidden_states = torch.cat(
            [
                torch.where(num_times_split > 0, accumulated_split / num_times_split, accumulated_split)
                for accumulated_split, num_times_split in zip(
                    accumulated_values.split(self.context_length, dim=1),
                    num_times_accumulated.split(self.context_length, dim=1),
                )
            ],
            dim=1,
        ).to(dtype)

        # 3. Feed-forward
        norm_hidden_states = self.norm3(hidden_states)

        if self._chunk_size is not None:
            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
        else:
            ff_output = self.ff(norm_hidden_states)

        hidden_states = ff_output + hidden_states
        if hidden_states.ndim == 4:
            hidden_states = hidden_states.squeeze(1)

        return hidden_states
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/attention.py