How to train adapter for Stable Diffusion 1.5 to teach it to interpret the T5 Encoder features?
I have this code, but it tries to train adapter to convert T5 embedings into CLIP L/CLIP G embedings, which is wrong how I think, cause the core idea of all this is to improve prompt following:
```py
# How to train adapter for SDXL, that it will accept T5 encoded text (tokenized and encoded by encoder) and will translate it to CLIP G and CLIP L tokens?

# Step 1. Import torch diffusers and transformers 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

import diffusers
from transformers import T5Tokenizer, T5EncoderModel
from datasets import load_dataset

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Step 2. Load the T5 tokenizer and encoder from Hugging Face Hub
t5_tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small", legacy=False, clean_up_tokenization_spaces=True)
t5_encoder: T5EncoderModel = T5EncoderModel.from_pretrained("google/flan-t5-small", device_map="auto", torch_dtype=torch.float16)

# Step 3. Load the SDXL model from Hugging Face Hub
sdxl = diffusers.StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16, 
    use_safetensors=True, 
    use_watermark=False,
    variant="fp16"
)

clip_l_tokenizer = sdxl.tokenizer
clip_l = sdxl.text_encoder.to(device)
clip_g_tokenizer = sdxl.tokenizer_2
clip_g = sdxl.text_encoder_2.to(device)

# Step 4. Write architecture of T5 Encoder Features -> CLIP adapter
# - Input: text (tokenized and encoded by encoder)
# - Output: CLIP G and CLIP L tokens
class T5ToCLIPAdapter(nn.Module):
    def __init__(self, input_dim, intermediate_dim, clip_l_dim, clip_g_dim, num_attention_heads, attention_dim, dropout=0.1):
        super().__init__()
        
        # Dense layers for feature transformation
        self.dense1 = nn.Linear(input_dim, intermediate_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        
        # Multi-Head Attention
        self.reduce_dim_layer = nn.Linear(intermediate_dim, attention_dim)
        self.multi_head_attention = nn.MultiheadAttention(embed_dim=attention_dim, num_heads=num_attention_heads)
        self.expand_dim_layer = nn.Linear(attention_dim, intermediate_dim)

        self.attention_norm = nn.LayerNorm(intermediate_dim)

        # CLIP Token Generation
        self.clip_l_projection = nn.Linear(intermediate_dim, clip_l_dim)
        self.clip_g_projection = nn.Linear(intermediate_dim, clip_g_dim)

    def forward(self, t5_features: torch.Tensor):
        # Feature transformation with dense layers and non-linear activation
        x = self.relu.forward(self.dense1.forward(t5_features))
        x = self.dropout.forward(x)
        
        # Reduce dimensions for attention processing
        x_reduced = self.reduce_dim_layer.forward(x)
        
        # Multi-Head Attention
        x_reduced = x_reduced.permute(1, 0, 2)  # Required permutation for PyTorch's attention module
        attention_output, _ = self.multi_head_attention.forward(x_reduced, x_reduced, x_reduced)
        attention_output = attention_output.permute(1, 0, 2)
        
        # Expand dimensions back to match the intermediate representation
        attention_output = self.expand_dim_layer.forward(attention_output)
        x = x + attention_output  # Add residual connection
        x = self.dropout.forward(self.attention_norm.forward(x))
        
        # Output CLIP Token projection
        clip_l_tokens = self.clip_l_projection.forward(x)
        clip_g_tokens = self.clip_g_projection.forward(x)

        return clip_l_tokens, clip_g_tokens
    
# Step 5: Define constants
INPUT_DIM = t5_encoder.config.d_model  # Assume T5 embedding dimension
INTERMEDIATE_DIM = 512  # Intermediate dimension for transformation
CLIP_L_DIM = 512  # Dimension for CLIP-L tokens
CLIP_G_DIM = 512  # Dimension for CLIP-G tokens
NUM_ATTENTION_HEADS = 8  # Number of attention heads
ATTENTION_DIM = 512  # Attention output dimension

# Step 6: Define dataset class
class TextToClipDataset(Dataset):
    def __init__(self, t5_features, clip_l_tokens, clip_g_tokens):
        self.t5_features = t5_features
        self.clip_l_tokens = clip_l_tokens
        self.clip_g_tokens = clip_g_tokens

    def __len__(self):
        return len(self.t5_features)

    def __getitem__(self, idx):
        return {
            't5_features': self.t5_features[idx],
            'clip_l_tokens': self.clip_l_tokens[idx],
            'clip_g_tokens': self.clip_g_tokens[idx]
        }

# Step 7: Train
def train_model(model: T5ToCLIPAdapter, dataloader: DataLoader, num_epochs: int, learning_rate=0.001):
    # Define loss function
    criterion = nn.MSELoss()

    # Define optimizer
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    model.train()  # Set model to training mode

    for epoch in range(num_epochs):
        running_loss = 0.0

        for batch in dataloader:
            t5_features = batch['t5_features']
            target_clip_l = batch['clip_l_tokens']
            target_clip_g = batch['clip_g_tokens']

            # Forward pass
            optimizer.zero_grad()  # Zero the gradients
            output_clip_l, output_clip_g = model(t5_features)

            # Compute loss
            loss_l = criterion(output_clip_l, target_clip_l)
            loss_g = criterion(output_clip_g, target_clip_g)
            loss = loss_l + loss_g

            # Backward pass and optimize
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * t5_features.size(0)

        epoch_loss = running_loss / len(dataloader.dataset)
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}')

    print('Training complete.')

t5_features = []
clip_l_tokens = []
clip_g_tokens = []

dataset_initial = load_dataset("OpenFace-CQUPT/HumanCaption-10M")
for entry_num in range(50):
    entry = dataset_initial['train'][entry_num]
    caption: str = entry['human_caption']
    
    t5_tokenized = t5_tokenizer(caption, return_tensors="pt", truncation=False).input_ids.to(device)
    t5_encoded = t5_encoder.forward(t5_tokenized).last_hidden_state
    t5_features.append(t5_encoded.to("cpu"))

    clip_l_tokenized = clip_l_tokenizer(caption, return_tensors="pt", truncation=False).input_ids.to(device)
    clip_l_shape_max_length = clip_l_tokenized.shape[-1]
    clip_l_encoded = []
    for i in range(0, clip_l_shape_max_length, clip_l_tokenizer.model_max_length):
        clip_l_encoded.append(clip_l.forward(clip_l_tokenized[:, i: i + clip_l_tokenizer.model_max_length])[0])
    clip_l_tokens.append(torch.cat(clip_l_encoded, dim=1))

    clip_g_tokenized = clip_g_tokenizer(caption, return_tensors="pt", truncation=False).input_ids.to(device)
    clip_g_shape_max_length = clip_g_tokenized.shape[-1]
    clip_g_encoded = []
    for i in range(0, clip_g_shape_max_length, clip_g_tokenizer.model_max_length):
        clip_g_encoded.append(clip_g.forward(clip_g_tokenized[:, i: i + clip_g_tokenizer.model_max_length])[0])
    clip_g_tokens.append(torch.cat(clip_g_encoded, dim=1))

dataset = TextToClipDataset(t5_features, clip_l_tokens, clip_g_tokens)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

model = T5ToCLIPAdapter(INPUT_DIM, INTERMEDIATE_DIM, CLIP_L_DIM, CLIP_G_DIM, NUM_ATTENTION_HEADS, ATTENTION_DIM).to(device)
train_model(model, dataloader, num_epochs=20)
```