Spaces:
Running on Zero
Running on Zero
| """AlignDRAW 512x512 | |
| Based on the supplied AlignDRAW PyTorch implementation. | |
| Image resolution: 512x512 | |
| DRAW generation steps: 69 | |
| Requires a newly trained checkpoint. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| def _lstm(pre, cell, H): | |
| i = torch.sigmoid(pre[:, 0:H]) | |
| f = torch.sigmoid(pre[:, H:2*H]) | |
| g = torch.tanh(pre[:, 2*H:3*H]) | |
| o = torch.sigmoid(pre[:, 3*H:4*H]) | |
| cell = f * cell + i * g | |
| return o * torch.tanh(cell), cell | |
| class SelectiveAttention: | |
| def __init__(self, H, W, N): | |
| self.H, self.W, self.N = H, W, N | |
| def params(self, raw): | |
| gx = (self.W + 1) / 2.0 * (raw[:, 0] + 1) | |
| gy = (self.H + 1) / 2.0 * (raw[:, 1] + 1) | |
| delta = (max(self.H, self.W) - 1.0) / (self.N - 1) * torch.exp(raw[:, 2]) | |
| sigma = torch.exp(raw[:, 3] / 2.0) | |
| gamma = torch.exp(raw[:, 4]).unsqueeze(1) | |
| return gx, gy, delta, sigma, gamma | |
| def _fb(self, g, delta, sigma, size): | |
| grid = torch.arange( | |
| self.N, device=g.device, dtype=g.dtype | |
| ) - self.N // 2 - 0.5 | |
| mu = g.unsqueeze(1) + delta.unsqueeze(1) * grid.unsqueeze(0) | |
| a = torch.arange(size, device=g.device, dtype=g.dtype) | |
| fb = torch.exp( | |
| -((a.view(1, 1, -1) - mu.unsqueeze(2)) ** 2) | |
| / (2 * sigma.view(-1, 1, 1) ** 2 + 1e-8) | |
| ) | |
| return fb / (fb.sum(-1, keepdim=True) + 1e-4) | |
| def read(self, img, gx, gy, delta, sigma): | |
| Fx = self._fb(gx, delta, sigma, self.W) | |
| Fy = self._fb(gy, delta, sigma, self.H) | |
| t = torch.einsum('bnh,bchw->bcnw', Fy, img) | |
| return torch.einsum('bcnw,bmw->bcnm', t, Fx) | |
| def write(self, win, gx, gy, delta, sigma): | |
| Fx = self._fb(gx, delta, sigma, self.W) | |
| Fy = self._fb(gy, delta, sigma, self.H) | |
| t = torch.einsum('bnh,bcnm->bchm', Fy, win) | |
| return torch.einsum('bchm,bmw->bchw', t, Fx) | |
| class AlignDRAW(nn.Module): | |
| def __init__( | |
| self, | |
| dimY=25323, | |
| dimLangRNN=128, | |
| dimAlign=512, | |
| imgSize=512, | |
| dimReadAttent=9, | |
| dimWriteAttent=9, | |
| dimRNNEnc=550, | |
| dimRNNDec=550, | |
| dimZ=275, | |
| runSteps=69 | |
| ): | |
| super().__init__() | |
| self.C = 3 | |
| self.H = self.W = imgSize | |
| self.dimX = self.C * self.H * self.W | |
| self.Nr, self.Nw = dimReadAttent, dimWriteAttent | |
| self.He, self.Hd = dimRNNEnc, dimRNNDec | |
| self.Z, self.T = dimZ, runSteps | |
| self.Hl, self.Al = dimLangRNN, dimAlign | |
| # Bidirectional language encoder | |
| self.emb_f = nn.Embedding(dimY, 4 * dimLangRNN) | |
| self.Wh_f = nn.Linear(dimLangRNN, 4 * dimLangRNN, bias=False) | |
| self.b_f = nn.Parameter(torch.zeros(4 * dimLangRNN)) | |
| self.emb_b = nn.Embedding(dimY, 4 * dimLangRNN) | |
| self.Wh_b = nn.Linear(dimLangRNN, 4 * dimLangRNN, bias=False) | |
| self.b_b = nn.Parameter(torch.zeros(4 * dimLangRNN)) | |
| # Caption alignment | |
| self.W_lang_align = nn.Linear( | |
| 2 * dimLangRNN, dimAlign, bias=False | |
| ) | |
| self.W_hdec_align = nn.Linear( | |
| dimRNNDec, dimAlign, bias=False | |
| ) | |
| self.b_align = nn.Parameter(torch.zeros(dimAlign)) | |
| self.v_align = nn.Parameter(torch.randn(dimAlign) * 0.01) | |
| self.W_s_hdec = nn.Linear( | |
| 2 * dimLangRNN, 4 * dimRNNDec, bias=False | |
| ) | |
| # Image reader | |
| self.W_read = nn.Linear(dimRNNDec, 5) | |
| read_dim = 2 * self.C * self.Nr * self.Nr | |
| # Encoder LSTM | |
| self.Wh_enc = nn.Linear( | |
| dimRNNEnc, 4 * dimRNNEnc, bias=False | |
| ) | |
| self.Wi_enc = nn.Linear( | |
| read_dim + dimRNNDec, | |
| 4 * dimRNNEnc, | |
| bias=False | |
| ) | |
| self.b_enc = nn.Parameter(torch.zeros(4 * dimRNNEnc)) | |
| self.W_mu = nn.Linear(dimRNNEnc, dimZ) | |
| self.W_logsigma = nn.Linear(dimRNNEnc, dimZ) | |
| # Decoder and learned prior | |
| self.Wh_dec = nn.Linear( | |
| dimRNNDec, 4 * dimRNNDec, bias=False | |
| ) | |
| self.Wz_dec = nn.Linear( | |
| dimZ, 4 * dimRNNDec, bias=False | |
| ) | |
| self.b_dec = nn.Parameter(torch.zeros(4 * dimRNNDec)) | |
| self.W_prior = nn.Linear(dimRNNDec, 2 * dimZ) | |
| # Image writer | |
| self.W_write = nn.Linear(dimRNNDec, 5) | |
| self.W_c = nn.Linear( | |
| dimRNNDec, self.C * self.Nw * self.Nw | |
| ) | |
| # Full 512x512 canvas initialization | |
| self.h0_lang = nn.Parameter(torch.zeros(1, dimLangRNN)) | |
| self.h0_enc = nn.Parameter(torch.zeros(1, dimRNNEnc)) | |
| self.h0_dec = nn.Parameter(torch.zeros(1, dimRNNDec)) | |
| self.c0 = nn.Parameter(torch.zeros(1, self.dimX)) | |
| self.reader = SelectiveAttention( | |
| self.H, self.W, self.Nr | |
| ) | |
| self.writer = SelectiveAttention( | |
| self.H, self.W, self.Nw | |
| ) | |
| self._init_weights() | |
| def _init_weights(self): | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Linear, nn.Embedding)): | |
| nn.init.normal_(m.weight, 0.0, 0.01) | |
| if getattr(m, "bias", None) is not None: | |
| nn.init.zeros_(m.bias) | |
| def encode_caption(self, y): | |
| b, L = y.shape | |
| yr = torch.flip(y, dims=[1]) | |
| hf = self.h0_lang.expand(b, -1).contiguous() | |
| hb = self.h0_lang.expand(b, -1).contiguous() | |
| cf = y.new_zeros((b, self.Hl), dtype=torch.float) | |
| cb = y.new_zeros((b, self.Hl), dtype=torch.float) | |
| outs_f, outs_b = [], [] | |
| for t in range(L): | |
| hf, cf = _lstm( | |
| self.emb_f(y[:, t]) + self.Wh_f(hf) + self.b_f, | |
| cf, self.Hl | |
| ) | |
| hb, cb = _lstm( | |
| self.emb_b(yr[:, t]) + self.Wh_b(hb) + self.b_b, | |
| cb, self.Hl | |
| ) | |
| outs_f.append(hf) | |
| outs_b.append(hb) | |
| hlang = torch.cat( | |
| [ | |
| torch.stack(outs_f, 1), | |
| torch.stack(outs_b, 1) | |
| ], | |
| dim=2 | |
| ) | |
| return hlang, self.W_lang_align(hlang) | |
| def _align(self, hlang, hid_align, h_dec): | |
| e = torch.tanh( | |
| hid_align | |
| + self.W_hdec_align(h_dec).unsqueeze(1) | |
| + self.b_align | |
| ) | |
| e = (e * self.v_align).sum(-1) | |
| alpha = F.softmax(e, dim=1) | |
| s = (alpha.unsqueeze(2) * hlang).sum(1) | |
| return s, alpha | |
| def forward(self, x, y): | |
| b = x.shape[0] | |
| assert x.shape[1] == self.dimX, ( | |
| f"Expected {self.dimX} image values, got {x.shape[1]}" | |
| ) | |
| h_dec = self.h0_dec.expand(b, -1).contiguous() | |
| cell_dec = x.new_zeros(b, self.Hd) | |
| h_enc = self.h0_enc.expand(b, -1).contiguous() | |
| cell_enc = x.new_zeros(b, self.He) | |
| c = self.c0.expand(b, -1).contiguous() | |
| mu_p = x.new_zeros(b, self.Z) | |
| ls_p = x.new_zeros(b, self.Z) | |
| hlang, hid_align = self.encode_caption(y) | |
| ximg = x.view(b, self.C, self.H, self.W) | |
| kl = x.new_zeros(()) | |
| for _ in range(self.T): | |
| x_hat = ximg - torch.sigmoid(c).view( | |
| b, self.C, self.H, self.W | |
| ) | |
| gx, gy, dl, sg, gm = self.reader.params( | |
| self.W_read(h_dec) | |
| ) | |
| r = gm[:, :, None, None] * torch.cat( | |
| [ | |
| self.reader.read(ximg, gx, gy, dl, sg), | |
| self.reader.read(x_hat, gx, gy, dl, sg) | |
| ], | |
| dim=1 | |
| ) | |
| r = r.reshape(b, -1) | |
| s, _ = self._align(hlang, hid_align, h_dec) | |
| h_enc, cell_enc = _lstm( | |
| self.Wh_enc(h_enc) | |
| + self.Wi_enc(torch.cat([r, h_dec], 1)) | |
| + self.b_enc, | |
| cell_enc, | |
| self.He | |
| ) | |
| mu = self.W_mu(h_enc) | |
| ls = 0.5 * self.W_logsigma(h_enc) | |
| z = mu + torch.exp(ls) * torch.randn_like(mu) | |
| kl = kl + torch.sum( | |
| -1 | |
| + ( | |
| (mu - mu_p) ** 2 + torch.exp(2 * ls) | |
| ) / torch.exp(2 * ls_p) | |
| - 2 * ls + 2 * ls_p | |
| ) | |
| h_dec, cell_dec = _lstm( | |
| self.Wh_dec(h_dec) | |
| + self.Wz_dec(z) | |
| + self.W_s_hdec(s) | |
| + self.b_dec, | |
| cell_dec, | |
| self.Hd | |
| ) | |
| mls = torch.tanh(self.W_prior(h_dec)) | |
| mu_p = mls[:, :self.Z] | |
| ls_p = mls[:, self.Z:] | |
| gx, gy, dl, sg, gm = self.writer.params( | |
| self.W_write(h_dec) | |
| ) | |
| win = self.W_c(h_dec).view( | |
| b, self.C, self.Nw, self.Nw | |
| ) | |
| c = c + ( | |
| 1.0 / gm[:, :, None, None] | |
| * self.writer.write(win, gx, gy, dl, sg) | |
| ).reshape(b, -1) | |
| c_final = torch.sigmoid(c) | |
| bce = F.binary_cross_entropy( | |
| c_final, x, reduction="sum" | |
| ) | |
| kl = 0.5 * kl | |
| return { | |
| "loss": (bce + kl) / b, | |
| "bce": bce / b, | |
| "kl": kl / b | |
| } | |
| def generate(self, y, steps=None, return_seq=False): | |
| steps = self.T if steps is None else steps | |
| b = y.shape[0] | |
| dev = self.h0_dec.device | |
| h_dec = self.h0_dec.expand(b, -1).contiguous() | |
| cell_dec = torch.zeros( | |
| b, self.Hd, device=dev | |
| ) | |
| c = self.c0.expand(b, -1).contiguous() | |
| mu_p = torch.zeros(b, self.Z, device=dev) | |
| ls_p = torch.zeros(b, self.Z, device=dev) | |
| hlang, hid_align = self.encode_caption(y) | |
| seq = [] | |
| for _ in range(steps): | |
| z = mu_p + torch.exp(ls_p) * torch.randn_like(mu_p) | |
| s, _ = self._align( | |
| hlang, hid_align, h_dec | |
| ) | |
| h_dec, cell_dec = _lstm( | |
| self.Wh_dec(h_dec) | |
| + self.Wz_dec(z) | |
| + self.W_s_hdec(s) | |
| + self.b_dec, | |
| cell_dec, | |
| self.Hd | |
| ) | |
| mls = torch.tanh(self.W_prior(h_dec)) | |
| mu_p = mls[:, :self.Z] | |
| ls_p = mls[:, self.Z:] | |
| gx, gy, dl, sg, gm = self.writer.params( | |
| self.W_write(h_dec) | |
| ) | |
| win = self.W_c(h_dec).view( | |
| b, self.C, self.Nw, self.Nw | |
| ) | |
| c = c + ( | |
| 1.0 / gm[:, :, None, None] | |
| * self.writer.write(win, gx, gy, dl, sg) | |
| ).reshape(b, -1) | |
| if return_seq: | |
| seq.append( | |
| torch.sigmoid(c).view( | |
| b, self.C, self.H, self.W | |
| ).cpu() | |
| ) | |
| img = torch.sigmoid(c).view( | |
| b, self.C, self.H, self.W | |
| ) | |
| return (img, seq) if return_seq else img | |