import torch
import torch.nn as nn

class SDXLResampler(nn.Module):
    def __init__(self, llm_dim=2048, sdxl_dim=2048, num_queries=77, depth=4, heads=16):
        super().__init__()
        
        # 1. The "Buckets" (77 fixed vectors that query the LLM output)
        self.latents = nn.Parameter(torch.randn(1, num_queries, sdxl_dim) * 0.02)
        
        # 2. Projector
        self.proj_in = nn.Linear(llm_dim, sdxl_dim) if llm_dim != sdxl_dim else nn.Identity()
        
        # 3. Attention Layers
        self.layers = nn.ModuleList([
            nn.TransformerDecoderLayer(
                d_model=sdxl_dim,
                nhead=heads,
                dim_feedforward=sdxl_dim * 4,
                dropout=0.0, 
                activation="gelu",
                batch_first=True,
                norm_first=True
            )
            for _ in range(depth)
        ])
        
        # 4. Output Norms
        self.norm_out = nn.LayerNorm(sdxl_dim)
        
        # 5. Pooled Head (For the "Vibe" vector [1280])
        self.pool_proj = nn.Linear(sdxl_dim, 1280)

    def forward(self, llm_hidden_states):
        """
        Input: [Batch, Seq_Len, LLM_Dim] (Qwen Output)
        Output: 
            - prompt_embeds: [Batch, 77, 2048]
            - pooled_embeds: [Batch, 1280]
        """
        batch_size = llm_hidden_states.shape[0]
        
        # Project Qwen dimension to SDXL dimension
        x = self.proj_in(llm_hidden_states)
        
        # Expand latents for the batch
        latents = self.latents.repeat(batch_size, 1, 1)
        
        # Apply Cross-Attention (Latents attend to Qwen Output)
        for layer in self.layers:
            latents = layer(tgt=latents, memory=x)
            
        latents = self.norm_out(latents)
        
        # Create Pooled Vector (Average of all 77 tokens -> projected)
        pooled = self.pool_proj(latents.mean(dim=1))
        
        return latents, pooled