59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
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
|