| """Character-level caption conditioning for cross-attention. |
| |
| Choice: we encode the WHOLE caption ("a navy sad face with 'DAILY' next to it"), max_len=64. |
| Longest possible caption is ~50 chars ("a purple smiley face with 'ABCDEFGH' next to it" = 48), |
| so 64 covers it with headroom and lets the model also handle free-form-ish prompts; the |
| template words are cheap tokens and the encoder learns to attend to the informative spans. |
| (Encoding only parsed parts would fit in 32 but bakes the parser into training.) |
| """ |
| import math |
| import torch, torch.nn as nn |
|
|
| MAX_LEN = 64 |
| VOCAB = 128 |
|
|
|
|
| def tokenize(texts, max_len=MAX_LEN): |
| """List[str] -> (ids LongTensor [B,L], mask BoolTensor [B,L], True = real token).""" |
| ids = torch.zeros(len(texts), max_len, dtype=torch.long) |
| for i, t in enumerate(texts): |
| b = [c if 0 < c < 128 else ord("?") for c in t.encode("ascii", "replace")][:max_len] |
| ids[i, :len(b)] = torch.tensor(b, dtype=torch.long) |
| return ids, ids != 0 |
|
|
|
|
| def sinusoidal_1d(L, d): |
| |
| pos = torch.arange(L).float()[:, None] |
| div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d)) |
| pe = torch.zeros(L, d); pe[:, 0::2] = torch.sin(pos * div); pe[:, 1::2] = torch.cos(pos * div) |
| return pe |
|
|
|
|
| class TextConditioner(nn.Module): |
| """Embeds chars -> per-token sequence [B,L,D] + mask [B,L]. No pooling. |
| |
| CFG dropout: with prob p_uncond a sample's whole sequence is replaced by a learned null |
| sequence (all positions marked valid so attention is never fully masked). |
| """ |
| def __init__(self, dim=256, layers=2, heads=4, max_len=MAX_LEN, p_uncond=0.1): |
| super().__init__() |
| self.max_len, self.p_uncond = max_len, p_uncond |
| self.emb = nn.Embedding(VOCAB, dim) |
| self.register_buffer("pe", sinusoidal_1d(max_len, dim), persistent=False) |
| layer = nn.TransformerEncoderLayer(dim, heads, dim * 4, dropout=0.0, batch_first=True, |
| norm_first=True, activation="gelu") |
| self.encoder = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False) |
| self.norm = nn.LayerNorm(dim) |
| self.null_seq = nn.Parameter(torch.randn(max_len, dim) * 0.02) |
|
|
| def forward(self, ids, mask, force_uncond=False): |
| x = self.emb(ids) + self.pe[None, :ids.shape[1]] |
| x = self.norm(self.encoder(x, src_key_padding_mask=~mask)) |
| B = ids.shape[0] |
| if force_uncond: |
| drop = torch.ones(B, dtype=torch.bool, device=ids.device) |
| elif self.training and self.p_uncond > 0: |
| drop = torch.rand(B, device=ids.device) < self.p_uncond |
| else: |
| return x, mask |
| null = self.null_seq[None, :ids.shape[1]].expand(B, -1, -1).to(x.dtype) |
| x = torch.where(drop[:, None, None], null, x) |
| mask = torch.where(drop[:, None], torch.ones_like(mask), mask) |
| return x, mask |
|
|
