update model.py

model.py

@@ -1,7 +1,9 @@
 
-"""AlignDRAW — checkpoint-compatible 32x32 PyTorch implementation.
-Original architecture dimensions preserved.
-Default generation steps: 69.
+"""AlignDRAW 512x512
+Based on the supplied AlignDRAW PyTorch implementation.
+Image resolution: 512x512
+DRAW generation steps: 69
+Requires a newly trained checkpoint.
 """
 
 import torch
@@ -10,10 +12,10 @@ import torch.nn.functional as F
 
 
 def _lstm(pre, cell, H):
-    i = torch.sigmoid(pre[:, 0 * H:1 * H])
-    f = torch.sigmoid(pre[:, 1 * H:2 * H])
-    g = torch.tanh(pre[:, 2 * H:3 * H])
-    o = torch.sigmoid(pre[:, 3 * H:4 * 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
 
@@ -31,35 +33,43 @@ class SelectiveAttention:
         return gx, gy, delta, sigma, gamma
 
     def _fb(self, g, delta, sigma, size):
-        grid = torch.arange(self.N, device=g.device).float() - self.N // 2 - 0.5
+        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).float()
+        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=32,
+        imgSize=512,
         dimReadAttent=9,
         dimWriteAttent=9,
         dimRNNEnc=550,
@@ -71,13 +81,14 @@ class AlignDRAW(nn.Module):
 
         self.C = 3
         self.H = self.W = imgSize
-        self.dimX = self.C * imgSize * imgSize
+        self.dimX = self.C * self.H * self.W
 
         self.Nr, self.Nw = dimReadAttent, dimWriteAttent
-        self.He, self.Hd, self.Z, self.T = dimRNNEnc, dimRNNDec, dimZ, runSteps
+        self.He, self.Hd = dimRNNEnc, dimRNNDec
+        self.Z, self.T = dimZ, runSteps
         self.Hl, self.Al = dimLangRNN, dimAlign
 
-        # Language encoder
+        # 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))
@@ -86,42 +97,75 @@ class AlignDRAW(nn.Module):
         self.Wh_b = nn.Linear(dimLangRNN, 4 * dimLangRNN, bias=False)
         self.b_b = nn.Parameter(torch.zeros(4 * dimLangRNN))
 
-        # Attention alignment
-        self.W_lang_align = nn.Linear(2 * dimLangRNN, dimAlign, bias=False)
-        self.W_hdec_align = nn.Linear(dimRNNDec, dimAlign, bias=False)
+        # 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)
 
-        # Reader and encoder
+        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
 
-        self.Wh_enc = nn.Linear(dimRNNEnc, 4 * dimRNNEnc, bias=False)
-        self.Wi_enc = nn.Linear(read_dim + dimRNNDec, 4 * dimRNNEnc, bias=False)
+        # 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.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)
 
-        # Writer
+        # Image writer
         self.W_write = nn.Linear(dimRNNDec, 5)
-        self.W_c = nn.Linear(dimRNNDec, self.C * self.Nw * self.Nw)
 
-        # Original checkpoint-compatible initial states
+        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.reader = SelectiveAttention(
+            self.H, self.W, self.Nr
+        )
+
+        self.writer = SelectiveAttention(
+            self.H, self.W, self.Nw
+        )
 
         self._init_weights()
 
@@ -129,6 +173,7 @@ class AlignDRAW(nn.Module):
         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)
 
@@ -149,15 +194,20 @@ class AlignDRAW(nn.Module):
                 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)],
+            [
+                torch.stack(outs_f, 1),
+                torch.stack(outs_b, 1)
+            ],
             dim=2
         )
 
@@ -165,16 +215,25 @@ class AlignDRAW(nn.Module):
 
     def _align(self, hlang, hid_align, h_dec):
         e = torch.tanh(
-            hid_align + self.W_hdec_align(h_dec).unsqueeze(1) + self.b_align
+            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)
 
@@ -187,21 +246,28 @@ class AlignDRAW(nn.Module):
         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))
+            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 = 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)
 
@@ -211,7 +277,8 @@ class AlignDRAW(nn.Module):
                 self.Wh_enc(h_enc)
                 + self.Wi_enc(torch.cat([r, h_dec], 1))
                 + self.b_enc,
-                cell_enc, self.He
+                cell_enc,
+                self.He
             )
 
             mu = self.W_mu(h_enc)
@@ -221,8 +288,9 @@ class AlignDRAW(nn.Module):
 
             kl = kl + torch.sum(
                 -1
-                + ((mu - mu_p) ** 2 + torch.exp(2 * ls))
-                / torch.exp(2 * ls_p)
+                + (
+                    (mu - mu_p) ** 2 + torch.exp(2 * ls)
+                ) / torch.exp(2 * ls_p)
                 - 2 * ls + 2 * ls_p
             )
 
@@ -231,13 +299,18 @@ class AlignDRAW(nn.Module):
                 + self.Wz_dec(z)
                 + self.W_s_hdec(s)
                 + self.b_dec,
-                cell_dec, self.Hd
+                cell_dec,
+                self.Hd
             )
 
             mls = torch.tanh(self.W_prior(h_dec))
-            mu_p, ls_p = mls[:, :self.Z], mls[:, self.Z:]
 
-            gx, gy, dl, sg, gm = self.writer.params(self.W_write(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
@@ -249,7 +322,11 @@ class AlignDRAW(nn.Module):
             ).reshape(b, -1)
 
         c_final = torch.sigmoid(c)
-        bce = F.binary_cross_entropy(c_final, x, reduction="sum")
+
+        bce = F.binary_cross_entropy(
+            c_final, x, reduction="sum"
+        )
+
         kl = 0.5 * kl
 
         return {
@@ -260,13 +337,16 @@ class AlignDRAW(nn.Module):
 
     @torch.no_grad()
     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)
+        cell_dec = torch.zeros(
+            b, self.Hd, device=dev
+        )
 
         c = self.c0.expand(b, -1).contiguous()
 
@@ -274,25 +354,34 @@ class AlignDRAW(nn.Module):
         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)
+            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
+                cell_dec,
+                self.Hd
             )
 
             mls = torch.tanh(self.W_prior(h_dec))
-            mu_p, ls_p = mls[:, :self.Z], mls[:, self.Z:]
 
-            gx, gy, dl, sg, gm = self.writer.params(self.W_write(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