def forward(self, x: torch.Tensor, latents: torch.Tensor, timestep_emb: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x (`torch.Tensor`):
                Image features.
            latents (`torch.Tensor`):
                Latent features.
            timestep_emb (`torch.Tensor`):
                Timestep embedding.

        Returns:
            `torch.Tensor`: Output latent features.
        """

        # Shift and scale for AdaLayerNorm
        emb = self.adaln_proj(self.adaln_silu(timestep_emb))
        shift_msa, scale_msa, shift_mlp, scale_mlp = emb.chunk(4, dim=1)

        # Fused Attention
        residual = latents
        x = self.ln0(x)
        latents = self.ln1(latents) * (1 + scale_msa[:, None]) + shift_msa[:, None]

        batch_size = latents.shape[0]

        query = self.attn.to_q(latents)
        kv_input = torch.cat((x, latents), dim=-2)
        key, value = self.attn.to_kv(kv_input).chunk(2, dim=-1)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // self.attn.heads

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

        weight = (query * self.attn.scale) @ (key * self.attn.scale).transpose(-2, -1)
        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
        latents = weight @ value

        latents = latents.transpose(1, 2).reshape(batch_size, -1, self.attn.heads * head_dim)
        latents = self.attn.to_out[0](latents)
        latents = self.attn.to_out[1](latents)
        latents = latents + residual

        ## FeedForward
        residual = latents
        latents = self.adaln_norm(latents) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
        return self.ff(latents) + residual
Forward pass.

        Args:
            x (`torch.Tensor`):
                Image features.
            latents (`torch.Tensor`):
                Latent features.
            timestep_emb (`torch.Tensor`):
                Timestep embedding.

        Returns:
            `torch.Tensor`: Output latent features.
        
src/diffusers/models/embeddings.py