update app.py

app.py

@@ -16,37 +16,54 @@ from model import AlignDRAW
 REPO = "mediasynthesismuseum/aligndraw"
 HERE = os.path.dirname(os.path.abspath(__file__))
 
-# Keep the original trained 32x32 model
-model = AlignDRAW()
-model.load_state_dict(
-    torch.load(
-        hf_hub_download(REPO, "aligndraw_coco.pt"),
-        map_location="cpu"
-    )
+# Original architecture dimensions, matching the original checkpoint.
+# 69 recurrent drawing steps.
+model = AlignDRAW(imgSize=32, runSteps=69)
+
+checkpoint_path = hf_hub_download(
+    REPO,
+    "aligndraw_coco.pt"
 )
-model.eval().to("cuda")
 
-DICT = pickle.load(
-    open(hf_hub_download(REPO, "dictionary.pkl"), "rb"),
-    encoding="latin1"
+checkpoint = torch.load(
+    checkpoint_path,
+    map_location="cpu"
 )
 
+model.load_state_dict(checkpoint)
+model.eval().to("cuda")
+
+# Original vocabulary
+with open(
+    hf_hub_download(REPO, "dictionary.pkl"),
+    "rb"
+) as f:
+    DICT = pickle.load(f, encoding="latin1")
+
 
 def _tokens(caption):
     words = caption.strip().lower().replace(".", " .").split()
+
     unknown = sorted({
-        w for w in words if w not in DICT and w != "."
+        w for w in words
+        if w not in DICT and w != "."
     })
+
     toks = [
-        DICT.get(w, DICT["UNK"]) for w in words
+        DICT.get(w, DICT["UNK"])
+        for w in words
     ] or [DICT["UNK"]]
+
     return toks, unknown
 
 
-# Convert original 32x32 output to 512x512
+# Enlarge original 32x32 output to 512x512
 def _up(t, size=512):
     a = (
-        t.detach().cpu().permute(1, 2, 0).numpy() * 255
+        t.detach()
+        .cpu()
+        .permute(1, 2, 0)
+        .numpy() * 255
     ).clip(0, 255).astype(np.uint8)
 
     return Image.fromarray(a).resize(
@@ -57,43 +74,59 @@ def _up(t, size=512):
 
 @spaces.GPU(duration=120)
 def generate(caption, n_samples):
-    if not caption.strip():
+
+    if not caption or not caption.strip():
         raise gr.Error("Type a caption first.")
 
     toks, unknown = _tokens(caption)
+
     y = torch.tensor(
-        [toks], dtype=torch.long, device="cuda"
+        [toks],
+        dtype=torch.long,
+        device="cuda"
     )
 
     with torch.inference_mode():
+
         imgs = model.generate(
             y.expand(int(n_samples), -1)
         )
 
         gallery = [
-            _up(imgs[i]) for i in range(int(n_samples))
+            _up(imgs[i])
+            for i in range(int(n_samples))
         ]
 
-        _, seq = model.generate(y, return_seq=True)
+        _, seq = model.generate(
+            y,
+            return_seq=True
+        )
 
-    # Preserve all original drawing frames.
+    # Original drawing sequence
     frames = [
-        np.array(_up(s[0])) for s in seq
+        np.array(_up(s[0]))
+        for s in seq
     ]
 
-    # Extend animation to 69 frames without changing the model.
+    # Keep the GIF at 69 displayed frames
     if frames:
         while len(frames) < 69:
             frames.append(frames[-1])
 
-    frames.extend([frames[-1]] * 8)
+        frames.extend([frames[-1]] * 8)
 
     gif = os.path.join(HERE, "draw.gif")
-    imageio.mimsave(gif, frames, duration=0.09, loop=0)
+
+    imageio.mimsave(
+        gif,
+        frames,
+        duration=0.09,
+        loop=0
+    )
 
     note = (
-        "💡 Not in the 2015 vocabulary (treated as UNK): "
-        + ", ".join(unknown)
+        "💡 Not in the 2015 vocabulary "
+        "(treated as UNK): " + ", ".join(unknown)
     ) if unknown else ""
 
     return gallery, gif, note
@@ -122,15 +155,16 @@ CSS = """
 """
 
 with gr.Blocks(title="AlignDRAW") as demo:
+
     gr.Markdown(
         "# 🖍️ AlignDRAW\n"
         "One of the **first text-to-image models** "
         "([Mansimov et al., 2015](https://arxiv.org/abs/1511.02793)) — "
         "resurrected.\n\n"
-        "Original COCO-trained AlignDRAW weights and vocabulary, "
-        "with generated images enlarged to **512×512**.\n\n"
-        "The original model still draws at 32×32 internally; "
-        "the larger display resolution does not add new image detail.\n\n"
+        "Original COCO-trained weights and vocabulary. "
+        "Images are displayed at **512×512**.\n\n"
+        "The original 32×32 model is preserved internally. "
+        "The drawing sequence uses 69 recurrent steps.\n\n"
         "Weights: [mediasynthesismuseum/aligndraw]"
         "(https://huggingface.co/mediasynthesismuseum/aligndraw).",
         elem_id="ad-title"
@@ -142,6 +176,7 @@ with gr.Blocks(title="AlignDRAW") as demo:
             value=EXAMPLES[0],
             scale=4
         )
+
         n = gr.Slider(
             4, 12, 8,
             step=1,
@@ -158,8 +193,9 @@ with gr.Blocks(title="AlignDRAW") as demo:
             columns=4,
             height=380
         )
+
         anim = gr.Image(
-            label="Drawing process (69 displayed frames)",
+            label="Drawing process (69 steps)",
             height=300
         )
 
@@ -169,7 +205,10 @@ with gr.Blocks(title="AlignDRAW") as demo:
         [gallery, anim, note]
     )
 
-    gr.Examples(EXAMPLES, inputs=cap)
+    gr.Examples(
+        EXAMPLES,
+        inputs=cap
+    )
 
 
 if __name__ == "__main__":