class LocalFacialExtractor(nn.Module):
    def __init__(
        self,
        id_dim: int = 1280,
        vit_dim: int = 1024,
        depth: int = 10,
        dim_head: int = 64,
        heads: int = 16,
        num_id_token: int = 5,
        num_queries: int = 32,
        output_dim: int = 2048,
        ff_mult: int = 4,
        num_scale: int = 5,
    ):
        super().__init__()

        # Storing identity token and query information
        self.num_id_token = num_id_token
        self.vit_dim = vit_dim
        self.num_queries = num_queries
        assert depth % num_scale == 0
        self.depth = depth // num_scale
        self.num_scale = num_scale
        scale = vit_dim**-0.5

        # Learnable latent query embeddings
        self.latents = nn.Parameter(torch.randn(1, num_queries, vit_dim) * scale)
        # Projection layer to map the latent output to the desired dimension
        self.proj_out = nn.Parameter(scale * torch.randn(vit_dim, output_dim))

        # Attention and ConsisIDFeedForward layer stack
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        PerceiverAttention(dim=vit_dim, dim_head=dim_head, heads=heads),  # Perceiver Attention layer
                        nn.Sequential(
                            nn.LayerNorm(vit_dim),
                            nn.Linear(vit_dim, vit_dim * ff_mult, bias=False),
                            nn.GELU(),
                            nn.Linear(vit_dim * ff_mult, vit_dim, bias=False),
                        ),  # ConsisIDFeedForward layer
                    ]
                )
            )

        # Mappings for each of the 5 different ViT features
        for i in range(num_scale):
            setattr(
                self,
                f"mapping_{i}",
                nn.Sequential(
                    nn.Linear(vit_dim, vit_dim),
                    nn.LayerNorm(vit_dim),
                    nn.LeakyReLU(),
                    nn.Linear(vit_dim, vit_dim),
                    nn.LayerNorm(vit_dim),
                    nn.LeakyReLU(),
                    nn.Linear(vit_dim, vit_dim),
                ),
            )

        # Mapping for identity embedding vectors
        self.id_embedding_mapping = nn.Sequential(
            nn.Linear(id_dim, vit_dim),
            nn.LayerNorm(vit_dim),
            nn.LeakyReLU(),
            nn.Linear(vit_dim, vit_dim),
            nn.LayerNorm(vit_dim),
            nn.LeakyReLU(),
            nn.Linear(vit_dim, vit_dim * num_id_token),
        )

    def forward(self, id_embeds: torch.Tensor, vit_hidden_states: List[torch.Tensor]) -> torch.Tensor:
        # Repeat latent queries for the batch size
        latents = self.latents.repeat(id_embeds.size(0), 1, 1)

        # Map the identity embedding to tokens
        id_embeds = self.id_embedding_mapping(id_embeds)
        id_embeds = id_embeds.reshape(-1, self.num_id_token, self.vit_dim)

        # Concatenate identity tokens with the latent queries
        latents = torch.cat((latents, id_embeds), dim=1)

        # Process each of the num_scale visual feature inputs
        for i in range(self.num_scale):
            vit_feature = getattr(self, f"mapping_{i}")(vit_hidden_states[i])
            ctx_feature = torch.cat((id_embeds, vit_feature), dim=1)

            # Pass through the PerceiverAttention and ConsisIDFeedForward layers
            for attn, ff in self.layers[i * self.depth : (i + 1) * self.depth]:
                latents = attn(ctx_feature, latents) + latents
                latents = ff(latents) + latents

        # Retain only the query latents
        latents = latents[:, : self.num_queries]
        # Project the latents to the output dimension
        latents = latents @ self.proj_out
        return latents
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/transformers/consisid_transformer_3d.py