Text Conditioning with CLIP and Cross-Attention

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

from from_noise_to_images import get_device
device = get_device()
print(f"Using device: {device}")

from from_noise_to_images.text_encoder import CLIPTextEncoder

# Load CLIP text encoder
text_encoder = CLIPTextEncoder(
    model_name="openai/clip-vit-base-patch32",
    device=device,
)

print(f"CLIP embedding dimension: {text_encoder.embed_dim}")
print(f"Max sequence length: {text_encoder.max_length}")

# Explore CLIP embeddings
prompts = [
    "a photo of a cat",
    "a photo of a dog",
    "a photo of a car",
    "a furry cat sitting",
    "a fluffy dog running",
]

# Get token embeddings and pooled embeddings
token_embeddings, pooled = text_encoder.encode(prompts, return_pooled=True)

print(f"Token embeddings shape: {token_embeddings.shape}")
print(f"  → (batch_size, max_length, embed_dim)")
print(f"Pooled embeddings shape: {pooled.shape}")
print(f"  → (batch_size, embed_dim)")

# Compute similarity matrix from pooled embeddings
pooled_norm = pooled / pooled.norm(dim=-1, keepdim=True)
similarity = (pooled_norm @ pooled_norm.T).cpu().numpy()

# Visualize
fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(similarity, cmap='viridis', vmin=0, vmax=1)
ax.set_xticks(range(len(prompts)))
ax.set_yticks(range(len(prompts)))
ax.set_xticklabels([p[:15] + "..." for p in prompts], rotation=45, ha='right')
ax.set_yticklabels([p[:15] + "..." for p in prompts])
ax.set_title('CLIP Embedding Similarity')
plt.colorbar(im)
plt.tight_layout()
plt.show()

print("\nKey Observations:")
print("  - 'cat' prompts are similar to each other")
print("  - 'dog' prompts are similar to each other")
print("  - Both are less similar to 'car'")

from from_noise_to_images.dit import CrossAttention

# Create a cross-attention layer
embed_dim = 256   # Image patch embedding dimension
context_dim = 512  # CLIP text embedding dimension
num_heads = 8

cross_attn = CrossAttention(
    embed_dim=embed_dim,
    context_dim=context_dim,
    num_heads=num_heads,
).to(device)

# Simulate inputs
batch_size = 2
num_patches = 64  # 8x8 grid
num_tokens = 10   # Text sequence length

# Random image patch embeddings
x = torch.randn(batch_size, num_patches, embed_dim, device=device)

# Random text embeddings (simulating CLIP output)
context = torch.randn(batch_size, num_tokens, context_dim, device=device)

# Forward pass
with torch.no_grad():
    output = cross_attn(x, context)

print(f"Cross-Attention Dimensions:")
print(f"  Image patches (X): {x.shape}")
print(f"  Text tokens (Z):   {context.shape}")
print(f"  Output:            {output.shape}")
print(f"\nThe output has the same shape as the image patches!")
print(f"Each patch now contains information from relevant text tokens.")

# Visualize attention weights for a real prompt
def visualize_cross_attention_weights(
    cross_attn: CrossAttention,
    x: torch.Tensor,
    context: torch.Tensor,
    tokens: list,
):
    """Visualize what each image patch attends to in the text."""
    B, N, _ = x.shape
    M = context.shape[1]
    H = cross_attn.num_heads
    d = cross_attn.head_dim
    
    with torch.no_grad():
        # Compute Q, K
        q = cross_attn.q_proj(x)
        k = cross_attn.k_proj(context)
        
        # Reshape for attention
        q = q.view(B, N, H, d).transpose(1, 2)
        k = k.view(B, M, H, d).transpose(1, 2)
        
        # Compute attention scores
        attn = (q @ k.transpose(-2, -1)) * cross_attn.scale
        attn = attn.softmax(dim=-1)  # (B, H, N, M)
    
    # Average over heads and batch
    attn_avg = attn[0].mean(dim=0).cpu().numpy()  # (N, M)
    
    # Reshape to 2D grid (assuming square patches)
    grid_size = int(np.sqrt(N))
    
    # Plot attention for each token
    num_tokens_to_show = min(6, M)
    fig, axes = plt.subplots(1, num_tokens_to_show, figsize=(3*num_tokens_to_show, 3))
    
    for i in range(num_tokens_to_show):
        attn_map = attn_avg[:, i].reshape(grid_size, grid_size)
        axes[i].imshow(attn_map, cmap='hot')
        token = tokens[i] if i < len(tokens) else "[PAD]"
        axes[i].set_title(f'"{token}"')
        axes[i].axis('off')
    
    plt.suptitle('Cross-Attention: Which patches attend to which tokens?', fontsize=12)
    plt.tight_layout()
    plt.show()

# Get real CLIP embeddings
test_prompt = "a red car driving fast"
test_embeddings, test_mask = text_encoder(test_prompt)
test_embeddings = test_embeddings.to(device)

# Get tokens for visualization
tokens = text_encoder.tokenizer.tokenize(test_prompt)
tokens = ['[BOS]'] + tokens + ['[EOS]']

# Random image patches (in practice, these would be from the DiT)
num_patches = 64
x_test = torch.randn(1, num_patches, embed_dim, device=device)

print(f"Prompt: '{test_prompt}'")
print(f"Tokens: {tokens}")
print()

visualize_cross_attention_weights(cross_attn, x_test, test_embeddings, tokens)

from from_noise_to_images.dit import TextConditionalDiT

# Create the text-conditional DiT for CIFAR-10 (32x32 RGB)
model = TextConditionalDiT(
    img_size=32,          # CIFAR-10 image size
    patch_size=4,         # 4x4 patches → 8x8 = 64 patches
    in_channels=3,        # RGB
    embed_dim=256,        # Patch embedding dimension
    depth=6,              # Number of transformer blocks
    num_heads=8,          # Attention heads
    mlp_ratio=4.0,        # MLP expansion
    context_dim=512,      # CLIP ViT-B/32 embedding dimension
).to(device)

# Count parameters
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"TextConditionalDiT Parameters: {num_params:,}")

# Compare to ConditionalDiT
from from_noise_to_images.dit import ConditionalDiT
class_model = ConditionalDiT(img_size=32, patch_size=4, in_channels=3)
class_params = sum(p.numel() for p in class_model.parameters() if p.requires_grad)
print(f"ConditionalDiT Parameters: {class_params:,}")
print(f"Difference: +{num_params - class_params:,} (cross-attention layers)")

# Test forward pass
batch_size = 4

# Random noisy images
x = torch.randn(batch_size, 3, 32, 32, device=device)

# Random timesteps
t = torch.rand(batch_size, device=device)

# Text embeddings from CLIP
prompts = [
    "a photo of a cat",
    "a photo of a dog",
    "a photo of a car",
    "a photo of a plane",
]
text_embeddings, text_mask = text_encoder(prompts)
text_embeddings = text_embeddings.to(device)
text_mask = text_mask.to(device)

# Forward pass
with torch.no_grad():
    v_pred = model(x, t, text_embeddings, text_mask)

print(f"Input shape:       {x.shape}")
print(f"Timesteps shape:   {t.shape}")
print(f"Text embed shape:  {text_embeddings.shape}")
print(f"Text mask shape:   {text_mask.shape}")
print(f"Output shape:      {v_pred.shape}")
print(f"\nModel correctly outputs velocity field with same shape as input!")

from from_noise_to_images.text_encoder import make_cifar10_captions

# Load CIFAR-10
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # [-1, 1]
])

train_dataset = torchvision.datasets.CIFAR10(
    root="./data",
    train=True,
    download=True,
    transform=transform
)

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=0,
    drop_last=True
)

# Test caption function
sample_labels = torch.tensor([0, 3, 5, 8])
captions = make_cifar10_captions(sample_labels)

print("Label → Caption mapping:")
for label, caption in zip(sample_labels.tolist(), captions):
    print(f"  {label} → '{caption}'")

# Visualize samples with their captions
images, labels = next(iter(train_loader))
captions = make_cifar10_captions(labels[:8])

fig, axes = plt.subplots(2, 4, figsize=(12, 6))

for i, (ax, img, caption) in enumerate(zip(axes.flat, images[:8], captions)):
    # Denormalize
    img = (img + 1) / 2
    img = img.permute(1, 2, 0).numpy()
    img = np.clip(img, 0, 1)
    
    ax.imshow(img)
    ax.set_title(caption, fontsize=9)
    ax.axis('off')

plt.suptitle('CIFAR-10 Samples with Text Captions', fontsize=12)
plt.tight_layout()
plt.show()

# Compare text embedding vs null embedding
_, text_emb = text_encoder.encode(["a photo of a cat"], return_pooled=True)
_, null_emb = text_encoder.encode([""], return_pooled=True)

# Compute similarity
text_norm = text_emb / text_emb.norm(dim=-1, keepdim=True)
null_norm = null_emb / null_emb.norm(dim=-1, keepdim=True)
similarity = (text_norm @ null_norm.mT).item()

print(f"Text embedding shape: {text_emb.shape}")
print(f"Null embedding shape: {null_emb.shape}")
print(f"Cosine similarity: {similarity:.3f}")
print(f"\nThe null embedding is different from any real text embedding.")
print(f"This makes CFG work: we can push away from 'nothing' toward 'something'.")

from from_noise_to_images.train import TextConditionalTrainer

# Create the trainer
trainer = TextConditionalTrainer(
    model=model,
    text_encoder=text_encoder,
    dataloader=train_loader,
    caption_fn=make_cifar10_captions,
    lr=1e-4,
    weight_decay=0.01,
    text_drop_prob=0.1,  # 10% text dropout for CFG
    device=device,
)

print("Training Text-Conditional DiT with CFG text dropout...")
print("(10% of samples trained with null text)\n")

NUM_EPOCHS = 30
losses = trainer.train(num_epochs=NUM_EPOCHS)

# Plot training loss
plt.figure(figsize=(10, 4))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Text-Conditional DiT Training Loss')
plt.grid(True, alpha=0.3)
plt.show()

print(f"Final loss: {losses[-1]:.4f}")

from from_noise_to_images.sampling import sample_text_conditional

model.eval()

# Generate images for each CIFAR-10 class
prompts = [
    "a photo of an airplane",
    "a photo of an automobile",
    "a photo of a bird",
    "a photo of a cat",
    "a photo of a deer",
    "a photo of a dog",
    "a photo of a frog",
    "a photo of a horse",
    "a photo of a ship",
    "a photo of a truck",
]

print("Generating images for each class with CFG scale=7.5...")

samples = sample_text_conditional(
    model=model,
    text_encoder=text_encoder,
    prompts=prompts,
    image_shape=(3, 32, 32),
    num_steps=50,
    cfg_scale=7.5,
    device=device,
)

# Display results
fig, axes = plt.subplots(2, 5, figsize=(15, 6))

for i, (ax, prompt) in enumerate(zip(axes.flat, prompts)):
    img = (samples[i] + 1) / 2  # Denormalize
    img = img.permute(1, 2, 0).cpu().numpy()
    img = np.clip(img, 0, 1)
    
    ax.imshow(img)
    ax.set_title(prompt.replace('a photo of ', ''), fontsize=10)
    ax.axis('off')

plt.suptitle('Text-to-Image Generation (CFG=7.5)', fontsize=14)
plt.tight_layout()
plt.show()

# Compare different CFG scales
test_prompt = "a photo of a cat"
cfg_scales = [1.0, 3.0, 5.0, 7.5, 10.0, 15.0]

print(f"Comparing CFG scales for: '{test_prompt}'")

fig, axes = plt.subplots(1, len(cfg_scales), figsize=(18, 3))

for ax, scale in zip(axes, cfg_scales):
    # Use same seed for fair comparison
    torch.manual_seed(42)
    
    sample = sample_text_conditional(
        model=model,
        text_encoder=text_encoder,
        prompts=[test_prompt],
        image_shape=(3, 32, 32),
        num_steps=50,
        cfg_scale=scale,
        device=device,
    )
    
    img = (sample[0] + 1) / 2
    img = img.permute(1, 2, 0).cpu().numpy()
    img = np.clip(img, 0, 1)
    
    ax.imshow(img)
    ax.set_title(f'scale={scale}')
    ax.axis('off')

plt.suptitle(f'Effect of CFG Scale: "{test_prompt}"', fontsize=12)
plt.tight_layout()
plt.show()

print("\nObservations:")
print("  scale=1.0: No guidance, blurry")
print("  scale=3-5: Mild guidance, improving")
print("  scale=7-10: Good balance")
print("  scale=15+: May oversaturate")

# Try some creative prompts
creative_prompts = [
    "a photo of a flying car",
    "a photo of a cat on a ship",
    "a photo of a red bird",
    "a photo of a fast truck",
]

print("Testing compositional understanding...")
print("(Results may be limited by CIFAR-10 training data)\n")

torch.manual_seed(123)
creative_samples = sample_text_conditional(
    model=model,
    text_encoder=text_encoder,
    prompts=creative_prompts,
    image_shape=(3, 32, 32),
    num_steps=50,
    cfg_scale=7.5,
    device=device,
)

fig, axes = plt.subplots(1, 4, figsize=(16, 4))

for ax, prompt, sample in zip(axes, creative_prompts, creative_samples):
    img = (sample + 1) / 2
    img = img.permute(1, 2, 0).cpu().numpy()
    img = np.clip(img, 0, 1)
    
    ax.imshow(img)
    ax.set_title(prompt, fontsize=10)
    ax.axis('off')

plt.suptitle('Compositional Prompts', fontsize=12)
plt.tight_layout()
plt.show()

# Save the trained model
trainer.save_checkpoint("phase4_text_conditional_dit.pt")
print("Model saved to phase4_text_conditional_dit.pt")