class Kandinsky3AttentionPooling(nn.Module):
    def __init__(self, num_channels, context_dim, head_dim=64):
        super().__init__()
        self.attention = Attention(
            context_dim,
            context_dim,
            dim_head=head_dim,
            out_dim=num_channels,
            out_bias=False,
        )

    def forward(self, x, context, context_mask=None):
        context_mask = context_mask.to(dtype=context.dtype)
        context = self.attention(context.mean(dim=1, keepdim=True), context, context_mask)
        return x + context.squeeze(1)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_kandinsky3.py