class PerceiverCrossAttention(nn.Module):
    def __init__(self, dim: int = 3072, dim_head: int = 128, heads: int = 16, kv_dim: int = 2048):
        super().__init__()

        self.scale = dim_head**-0.5
        self.dim_head = dim_head
        self.heads = heads
        inner_dim = dim_head * heads

        # Layer normalization to stabilize training
        self.norm1 = nn.LayerNorm(dim if kv_dim is None else kv_dim)
        self.norm2 = nn.LayerNorm(dim)

        # Linear transformations to produce queries, keys, and values
        self.to_q = nn.Linear(dim, inner_dim, bias=False)
        self.to_kv = nn.Linear(dim if kv_dim is None else kv_dim, inner_dim * 2, bias=False)
        self.to_out = nn.Linear(inner_dim, dim, bias=False)

    def forward(self, image_embeds: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
        # Apply layer normalization to the input image and latent features
        image_embeds = self.norm1(image_embeds)
        hidden_states = self.norm2(hidden_states)

        batch_size, seq_len, _ = hidden_states.shape

        # Compute queries, keys, and values
        query = self.to_q(hidden_states)
        key, value = self.to_kv(image_embeds).chunk(2, dim=-1)

        # Reshape tensors to split into attention heads
        query = query.reshape(query.size(0), -1, self.heads, self.dim_head).transpose(1, 2)
        key = key.reshape(key.size(0), -1, self.heads, self.dim_head).transpose(1, 2)
        value = value.reshape(value.size(0), -1, self.heads, self.dim_head).transpose(1, 2)

        # Compute attention weights
        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
        weight = (query * scale) @ (key * scale).transpose(-2, -1)  # More stable scaling than post-division
        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)

        # Compute the output via weighted combination of values
        out = weight @ value

        # Reshape and permute to prepare for final linear transformation
        out = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, -1)

        return self.to_out(out)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/transformers/consisid_transformer_3d.py