def calculate_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn: Attention,
        batch_size: int,
        image_seq_length: int,
        text_seq_length: int,
        attention_mask: Optional[torch.Tensor],
        image_rotary_emb: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""
        Core attention computation with inversion-guided RoPE integration.

        Args:
            query (`torch.Tensor`): `[batch_size, seq_len, dim]` query tensor
            key (`torch.Tensor`): `[batch_size, seq_len, dim]` key tensor
            value (`torch.Tensor`): `[batch_size, seq_len, dim]` value tensor
            attn (`Attention`): Parent attention module with projection layers
            batch_size (`int`): Effective batch size (after chunk splitting)
            image_seq_length (`int`): Length of image feature sequence
            text_seq_length (`int`): Length of text feature sequence
            attention_mask (`Optional[torch.Tensor]`): Attention mask tensor
            image_rotary_emb (`Optional[torch.Tensor]`): Rotary embeddings for image positions

        Returns:
            `Tuple[torch.Tensor, torch.Tensor]`:
                (1) hidden_states: [batch_size, image_seq_length, dim] processed image features
                (2) encoder_hidden_states: [batch_size, text_seq_length, dim] processed text features
        """
        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads

        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        # Apply RoPE if needed
        if image_rotary_emb is not None:
            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
            if not attn.is_cross_attention:
                if key.size(2) == query.size(2):  # Attention for reference hidden states
                    key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
                else:  # RoPE should be applied to each group of image tokens
                    key[:, :, text_seq_length : text_seq_length + image_seq_length] = apply_rotary_emb(
                        key[:, :, text_seq_length : text_seq_length + image_seq_length], image_rotary_emb
                    )
                    key[:, :, text_seq_length * 2 + image_seq_length :] = apply_rotary_emb(
                        key[:, :, text_seq_length * 2 + image_seq_length :], image_rotary_emb
                    )

        hidden_states = F.scaled_dot_product_attention(
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        encoder_hidden_states, hidden_states = hidden_states.split(
            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
        )
        return hidden_states, encoder_hidden_states

        Core attention computation with inversion-guided RoPE integration.

        Args:
            query (`torch.Tensor`): `[batch_size, seq_len, dim]` query tensor
            key (`torch.Tensor`): `[batch_size, seq_len, dim]` key tensor
            value (`torch.Tensor`): `[batch_size, seq_len, dim]` value tensor
            attn (`Attention`): Parent attention module with projection layers
            batch_size (`int`): Effective batch size (after chunk splitting)
            image_seq_length (`int`): Length of image feature sequence
            text_seq_length (`int`): Length of text feature sequence
            attention_mask (`Optional[torch.Tensor]`): Attention mask tensor
            image_rotary_emb (`Optional[torch.Tensor]`): Rotary embeddings for image positions

        Returns:
            `Tuple[torch.Tensor, torch.Tensor]`:
                (1) hidden_states: [batch_size, image_seq_length, dim] processed image features
                (2) encoder_hidden_states: [batch_size, text_seq_length, dim] processed text features
        
examples/community/cogvideox_ddim_inversion.py