Async pre-generation, fix loader, fix collection display

Key changes:
- Pre-generate reward image immediately on "Let's go!" so it's ready when user earns it
- locked placeholder appears in collection the moment generation starts
- Image unlocks (reveal animation) when user earns the reward
- Fallback on-demand generation if pre-image not ready in time
- Fix loader: use value-swapping ("" vs LOADER_HTML) instead of visibility toggle — more reliable on HF Spaces
- Fix collection: drop canvas CORS approach, use Python→JS base64 injection via gr.HTML polling
- Collection stores base64 JPEGs in localStorage, supports tap-to-zoom modal
- Event chain: go_btn → start_game → pregenerate; check_btn → check_answer → generate_on_demand → pregenerate

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

CLAUDE.md

@@ -36,10 +36,30 @@ Full loop shortcut: `bash scripts/dev.sh` (runs steps 1–4, then run test.sh ma
 
 **Game state dict keys:**
 `name`, `level`, `score`, `streak`, `correct_since_reward`, `level_correct`,
-`question`, `answer`, `selected_animals`, `selected_places`, `generate_now`
+`question`, `answer`, `selected_animals`, `selected_places`,
+`reward_count`, `pregenerate_next`, `generate_now`, `pending_reward_id`
 
 **Reward trigger:** every `REWARD_EVERY = 3` correct answers.
-Image generated via `.then()` chaining on the Check button — synchronous, within Gradio request context (required for ZeroGPU).
+
+**Pre-generation flow (async image UX):**
+1. "Let's go!" → `start_game()` → adds locked placeholder for reward #1 in collection →
+   `.then(pregenerate_image)` starts background generation immediately
+2. `pregenerate_image` stores PIL image in `hidden_image` (gr.Image) and base64 data URL in
+   `hidden_data_url` (gr.Textbox), sets `state["pending_reward_id"]`
+3. When user earns a reward: `check_answer` receives `hidden_image` + `hidden_data_url` as inputs.
+   - If pre-image matches reward_id → show immediately + unlock in collection + kick off next pregenerate
+   - If not ready → show LOADER_HTML + set `generate_now=True`
+4. `.then(generate_on_demand)` → generates on-demand if `generate_now`, unlocks collection
+5. `.then(pregenerate_image)` → generates next reward in background
+
+**Collection system:**
+- `collection_trigger = gr.HTML("")` receives JSON-encoded action payloads
+- JS polls `#numzoo-coll-data` element every 200ms for `{ts, actions:[{action,id,src}]}`
+- Actions: `add-locked` (locked placeholder), `unlock` (reveal image), `failed` (remove)
+- Images stored in localStorage as base64 JPEGs (key: `numzoo_collection_v2`)
+
+**Loader fix:** Use value-swapping (`""` = hidden, `LOADER_HTML` = shown) instead of `visible` flag.
+This is more reliable in Gradio 6 than toggling visibility on gr.HTML components.
 
 ## Gradio rules
 - Use `gr.update(visible=True/False)` for visibility — never return bare booleans

app.py

@@ -2,8 +2,13 @@
 🦁 NumZoo — Math practice with cute AI-generated animal rewards!
 """
 
+import base64
+import io
+import json
 import random
+
 import gradio as gr
+
 from math_engine import generate_question, LEVEL_NAMES, LEVEL_THRESHOLDS, level_up_message
 from image_generator import generate_reward_image
 
@@ -28,6 +33,21 @@ def _status_text(state: dict) -> str:
 def _safe_question(state: dict) -> str:
     return f"## {state.get('question', '')} = ?"
 
+def _img_to_data_url(pil_image) -> str:
+    """Convert PIL image to JPEG data URL for localStorage."""
+    buf = io.BytesIO()
+    pil_image.save(buf, format="JPEG", quality=80)
+    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
+
+def _coll(*actions) -> str:
+    """Encode collection actions into the collection_trigger HTML component.
+    Each action is a (action_name, reward_id, src) tuple."""
+    payload = json.dumps({
+        "ts": random.random(),
+        "actions": [{"action": a, "id": str(i), "src": s} for a, i, s in actions],
+    })
+    return f"<div id='numzoo-coll-data' data-payload='{payload}'></div>"
+
 # ---------------------------------------------------------------------------
 # Step 1 — Name entry
 # ---------------------------------------------------------------------------
@@ -56,45 +76,89 @@ def start_game(animal_sel: list, place_sel: list, state: dict):
     try:
         if not animal_sel or not place_sel:
             return (state, gr.update(visible=True), gr.update(visible=False),
-                    "⚠️ Pick at least one animal and one place!", "", "")
+                    "⚠️ Pick at least one animal and one place!", "", "", "")
 
         state.update({
-            "selected_animals":       animal_sel[:3],
-            "selected_places":        place_sel[:3],
-            "level":                  1,
-            "score":                  0,
-            "streak":                 0,
-            "correct_since_reward":   0,
-            "level_correct":          0,
+            "selected_animals":     animal_sel[:3],
+            "selected_places":      place_sel[:3],
+            "level":                1,
+            "score":                0,
+            "streak":               0,
+            "correct_since_reward": 0,
+            "level_correct":        0,
+            "reward_count":         0,
+            "pregenerate_next":     True,   # trigger background generation immediately
+            "generate_now":         False,
         })
 
         question, answer = generate_question(state["level"])
         state["question"] = question
         state["answer"]   = answer
 
+        # Add locked placeholder for reward #1 immediately
         return (state, gr.update(visible=False), gr.update(visible=True),
-                "", _status_text(state), _safe_question(state))
+                "", _status_text(state), _safe_question(state),
+                _coll(("add-locked", 1, "")))
 
     except Exception as e:
         print(f"start_game error: {e}")
         return (state, gr.update(visible=True), gr.update(visible=False),
-                "⚠️ Something went wrong, try again.", "", "")
+                "⚠️ Something went wrong, try again.", "", "", "")
+
+# ---------------------------------------------------------------------------
+# Background pre-generation — starts right after "Let's go!" and again after
+# each reward, so the next image is ready before the user earns it.
+# ---------------------------------------------------------------------------
+
+def pregenerate_image(state: dict):
+    """Generate next reward image in background. No-op if pregenerate_next is False."""
+    if not state.get("pregenerate_next"):
+        return state, gr.update(), gr.update(), ""
+
+    state["pregenerate_next"] = False
+    next_id = state.get("reward_count", 0) + 1
+
+    try:
+        animals = state.get("selected_animals", [random.choice(ANIMAL_EMOJIS)])
+        places  = state.get("selected_places",  [random.choice(PLACE_EMOJIS)])
+        streak  = state.get("streak", 0)
+
+        print(f"[pregenerate] generating reward #{next_id}")
+        result, prompt = generate_reward_image(streak, animals, places)
+        print(f"[pregenerate] reward #{next_id} done | prompt={prompt!r}")
+
+        if result is not None:
+            data_url = _img_to_data_url(result)
+            state["pending_reward_id"] = next_id
+            # Store in hidden components — collection stays LOCKED until user earns reward
+            return state, result, data_url, ""
+
+        state.pop("pending_reward_id", None)
+        return state, None, "", ""
+
+    except Exception as e:
+        import traceback
+        print(f"[pregenerate] ❌ {e}\n{traceback.format_exc()}")
+        state.pop("pending_reward_id", None)
+        return state, None, "", ""
 
 # ---------------------------------------------------------------------------
-# Step 3 — Answer checking (no reward yet, just feedback)
+# Step 3 — Answer checking
 # ---------------------------------------------------------------------------
 
-def check_answer(user_input: str, state: dict):
+def check_answer(user_input: str, state: dict, pre_image, pre_data_url: str):
     try:
         if not state or not state.get("question"):
             return (state, _status_text(state), _safe_question(state), "",
-                    "", gr.update(visible=False), gr.update(visible=False), None, "")
+                    "", gr.update(visible=False), "", gr.update(visible=False), "",
+                    gr.update(), gr.update(), "")
 
         try:
             user_answer = int(user_input.strip())
         except (ValueError, AttributeError):
             return (state, _status_text(state), _safe_question(state), "",
-                    "⚠️ Numbers only!", gr.update(visible=False), gr.update(visible=False), None, "")
+                    "⚠️ Numbers only!", gr.update(visible=False), "", gr.update(visible=False), "",
+                    gr.update(), gr.update(), "")
 
         correct = (user_answer == state["answer"])
 
@@ -104,7 +168,6 @@ def check_answer(user_input: str, state: dict):
             state["correct_since_reward"] += 1
             state["level_correct"]        += 1
 
-            # Level-up?
             threshold = LEVEL_THRESHOLDS.get(state["level"], 999)
             level_msg = ""
             if state["level_correct"] >= threshold and state["level"] < 4:
@@ -115,19 +178,47 @@ def check_answer(user_input: str, state: dict):
             streak_fire = "🔥" * min(state["streak"], 5)
             feedback    = f"✅ {level_msg}" if level_msg else f"✅ Great! {streak_fire}"
 
-            # Reward due? → show loading panel, trigger generate step
             if state["correct_since_reward"] >= REWARD_EVERY:
                 state["correct_since_reward"] = 0
-                state["generate_now"] = True       # signal .then() chain to generate
+                state["reward_count"] = state.get("reward_count", 0) + 1
+                reward_id = state["reward_count"]
+
                 question, answer = generate_question(state["level"])
                 state["question"] = question
                 state["answer"]   = answer
-                return (state, _status_text(state), _safe_question(state), "",
-                        feedback,
-                        gr.update(visible=True),    # reward_panel visible
-                        gr.update(visible=True),    # loader visible
-                        gr.update(visible=False),   # image hidden until ready
-                        "")                         # no error
+
+                # Check if a pre-generated image is ready for this exact reward
+                pending_id = state.get("pending_reward_id")
+                has_pre    = pre_image is not None and bool(pre_data_url) and pending_id == reward_id
+
+                if has_pre:
+                    # ✅ Pre-generated image ready — show immediately, unlock in collection
+                    state.pop("pending_reward_id", None)
+                    state["pregenerate_next"] = True      # kick off next background gen
+                    next_id = reward_id + 1
+                    return (state, _status_text(state), _safe_question(state), "",
+                            feedback,
+                            gr.update(visible=True),                      # reward_panel
+                            "",                                           # loader cleared
+                            gr.update(visible=True, value=pre_image),     # image shown
+                            "",
+                            None,                                         # clear hidden_image
+                            "",                                           # clear hidden_data_url
+                            _coll(("unlock",    reward_id, pre_data_url),
+                                  ("add-locked", next_id,   "")))
+                else:
+                    # ⏳ On-demand generation (image not ready yet)
+                    state["generate_now"]     = True
+                    state["pregenerate_next"] = False
+                    return (state, _status_text(state), _safe_question(state), "",
+                            feedback,
+                            gr.update(visible=True),    # reward_panel
+                            LOADER_HTML,                # show loader
+                            gr.update(visible=False),   # image hidden
+                            "",
+                            gr.update(),               # keep hidden_image
+                            gr.update(),               # keep hidden_data_url
+                            "")                         # collection stays locked
 
         else:
             state["streak"] = 0
@@ -140,55 +231,58 @@ def check_answer(user_input: str, state: dict):
 
         return (state, _status_text(state), _safe_question(state), "",
                 feedback,
-                gr.update(visible=False),  # reward_panel
-                gr.update(visible=False),  # loader
-                None,
-                "")
+                gr.update(visible=False), "", gr.update(visible=False), "",
+                gr.update(), gr.update(), "")
 
     except Exception as e:
-        print(f"check_answer error: {e}")
+        import traceback
+        print(f"check_answer error: {e}\n{traceback.format_exc()}")
         state["generate_now"] = False
         question, answer = generate_question(state.get("level", 1))
         state["question"] = question
         state["answer"]   = answer
         return (state, _status_text(state), _safe_question(state), "",
                 "⚠️ Something went wrong!",
-                gr.update(visible=False), gr.update(visible=False), None, "")
+                gr.update(visible=False), "", gr.update(visible=False), "",
+                gr.update(), gr.update(), "")
 
 # ---------------------------------------------------------------------------
-# Step 4 — Actually generate the image (called by gr.Timer after reward shown)
+# On-demand generation — fallback when pre-image wasn't ready
 # ---------------------------------------------------------------------------
 
-def generate_image(state: dict):
-    """Chained generation — only runs if check_answer flagged generate_now."""
+def generate_on_demand(state: dict):
+    """Generate image on demand. No-op if generate_now is False."""
     if not state.get("generate_now"):
-        return gr.update(), gr.update(), gr.update()   # no-op
+        return state, "", gr.update(), "", ""
 
-    state["generate_now"] = False
+    state["generate_now"]     = False
+    state["pregenerate_next"] = True   # kick off next background gen after this
+    reward_id = state.get("reward_count", 1)
 
     try:
         animals = state.get("selected_animals", [random.choice(ANIMAL_EMOJIS)])
         places  = state.get("selected_places",  [random.choice(PLACE_EMOJIS)])
         streak  = state.get("streak", 0)
 
-        print(f"[generate_image] calling generate_reward_image | streak={streak} | animals={animals} | places={places}")
+        print(f"[on_demand] generating reward #{reward_id}")
         result, prompt = generate_reward_image(streak, animals, places)
-        print(f"[generate_image] result={result} | prompt={prompt!r}")
+        print(f"[on_demand] done | reward #{reward_id}")
 
-        if result is None:
-            print("[generate_image] ⚠️  result is None — generation failed silently inside image_generator")
-            return (gr.update(visible=False),
-                    gr.update(visible=False),
-                    "⚠️ Could not generate image — try again later!")
+        if result is not None:
+            data_url = _img_to_data_url(result)
+            next_id  = reward_id + 1
+            return (state, "",
+                    gr.update(visible=True, value=result),
+                    "",
+                    _coll(("unlock",    reward_id, data_url),
+                          ("add-locked", next_id,   "")))
 
-        print("[generate_image] ✅ image generated successfully")
-        return gr.update(visible=False), gr.update(visible=True, value=result), ""
+        return state, "", gr.update(visible=False), "⚠️ Could not generate image — try again later!", ""
 
     except Exception as e:
         import traceback
-        print(f"[generate_image] ❌ exception: {e}")
-        print(traceback.format_exc())
-        return gr.update(visible=False), gr.update(visible=False), f"⚠️ Error: {e}"
+        print(f"[on_demand] ❌ {e}\n{traceback.format_exc()}")
+        return state, "", gr.update(visible=False), f"⚠️ Error: {e}", ""
 
 # ---------------------------------------------------------------------------
 # Restart
@@ -196,7 +290,7 @@ def generate_image(state: dict):
 
 def restart(state: dict):
     return ({}, gr.update(visible=True), gr.update(visible=False),
-            gr.update(visible=False), gr.update(), "")
+            gr.update(visible=False), gr.update(), "", None, "")
 
 # ---------------------------------------------------------------------------
 # Gradio UI
@@ -212,7 +306,6 @@ CSS = """
 #feedback-box { text-align: center; font-size: 1.3em; min-height: 2em; }
 #status-box   { text-align: center; padding: 0.4em; border-radius: 8px; }
 #reward-img   { border-radius: 16px; }
-#loader       { text-align: center; font-size: 1.4em; padding: 2em; }
 .answer-input input { font-size: 2em !important; text-align: center !important; }
 /* Emoji picker grid — animals: 6 cols × 2 rows, places: 5 cols × 2 rows */
 .emoji-group .wrap {
@@ -224,38 +317,141 @@ CSS = """
 #animal-picker .wrap { grid-template-columns: repeat(6, 52px) !important; }
 #place-picker  .wrap { grid-template-columns: repeat(5, 52px) !important; }
 .emoji-group label {
-  width: 52px !important;
-  height: 52px !important;
-  display: flex !important;
-  align-items: center !important;
-  justify-content: center !important;
-  font-size: 1.8em !important;
-  cursor: pointer !important;
-  border-radius: 12px !important;
-  border: 2px solid transparent !important;
-  transition: all 0.15s !important;
-  user-select: none !important;
+  width: 52px !important; height: 52px !important;
+  display: flex !important; align-items: center !important; justify-content: center !important;
+  font-size: 1.8em !important; cursor: pointer !important;
+  border-radius: 12px !important; border: 2px solid transparent !important;
+  transition: all 0.15s !important; user-select: none !important;
 }
 .emoji-group label:has(input:checked) {
-  background: #e9d5ff !important;
-  border-color: #7c3aed !important;
+  background: #e9d5ff !important; border-color: #7c3aed !important;
 }
 .emoji-group input[type="checkbox"] { display: none !important; }
+/* Unlock reveal animation */
+@keyframes numzoo-unblur {
+  from { filter: blur(16px) brightness(0.4); transform: scale(0.95); }
+  to   { filter: blur(0)    brightness(1);   transform: scale(1); }
+}
+.numzoo-unlock { animation: numzoo-unblur 0.7s ease-out forwards; }
+/* Collection locked pulse */
+@keyframes numzoo-pulse { 0%,100% { opacity:.4; } 50% { opacity:.9; } }
 """
 
 LOADER_HTML = """
-<div id="loader" style="text-align:center; padding:2em;">
-  <div style="font-size:3em; animation: spin 1s linear infinite; display:inline-block;">🎨</div>
+<div style="text-align:center; padding:2em;">
+  <div style="font-size:3em; animation:spin 1s linear infinite; display:inline-block;">🎨</div>
   <div style="margin-top:0.5em; color:#888; font-size:1.1em;">Generating your reward…</div>
 </div>
-<style>
-  @keyframes spin { 0%{transform:rotate(0deg)} 100%{transform:rotate(360deg)} }
-</style>
+<style>@keyframes spin{0%{transform:rotate(0deg)}100%{transform:rotate(360deg)}}</style>
+"""
+
+COLLECTION_HTML = """
+<div id="numzoo-collection-wrap" style="margin-top:16px;">
+  <div id="numzoo-collection-header"
+       style="display:none; font-weight:bold; font-size:1.05em; margin:0 0 8px 4px;">🖼️ Collection</div>
+  <div id="numzoo-collection"
+       style="display:flex; flex-wrap:wrap; gap:8px; justify-content:center;"></div>
+</div>
+
+<!-- Zoom modal -->
+<div id="numzoo-modal"
+     style="display:none; position:fixed; inset:0; background:rgba(0,0,0,0.88);
+            z-index:9999; align-items:center; justify-content:center; cursor:pointer;">
+  <img id="numzoo-modal-img"
+       style="max-width:92vw; max-height:92vh; border-radius:20px; box-shadow:0 8px 40px #0006;">
+</div>
+
+<script>
+(function() {
+  var KEY = 'numzoo_collection_v2', MAX = 30;
+
+  function getCol() { try { return JSON.parse(localStorage.getItem(KEY)||'[]'); } catch(e) { return []; } }
+  function saveCol(c) { localStorage.setItem(KEY, JSON.stringify(c.slice(-MAX))); }
+
+  window.numzooAddLocked = function(id) {
+    id = String(id);
+    var col = getCol();
+    if (!col.find(function(x){return x.id===id;})) {
+      col.push({id:id, status:'locked', src:null});
+      saveCol(col);
+    }
+    renderCollection();
+  };
+
+  window.numzooUnlock = function(id, src) {
+    id = String(id);
+    var col = getCol(), item = col.find(function(x){return x.id===id;});
+    if (item) { item.status='unlocked'; item.src=src; }
+    else { col.push({id:id, status:'unlocked', src:src}); }
+    saveCol(col);
+    renderCollection();
+  };
+
+  window.numzooMarkFailed = function(id) {
+    id = String(id);
+    saveCol(getCol().filter(function(x){return x.id!==id;}));
+    renderCollection();
+  };
+
+  function renderCollection() {
+    var col = getCol(), wrap = document.getElementById('numzoo-collection'),
+        hdr = document.getElementById('numzoo-collection-header');
+    if (!wrap) return;
+    wrap.innerHTML = '';
+    if (!col.length) { hdr.style.display='none'; return; }
+    hdr.style.display = 'block';
+    col.slice().reverse().forEach(function(item) {
+      var el = document.createElement('div');
+      el.style.cssText = 'position:relative;width:80px;height:80px;border-radius:12px;overflow:hidden;flex-shrink:0;';
+      if (item.status === 'locked') {
+        el.style.background = '#1e1e2e';
+        el.style.border = '2px dashed #555';
+        el.innerHTML =
+          '<div style="width:100%;height:100%;display:flex;align-items:center;justify-content:center;font-size:2em;z-index:1;position:relative;">🔒</div>' +
+          '<div style="position:absolute;inset:0;background:linear-gradient(135deg,#2d2d44,#111128);animation:numzoo-pulse 2s ease-in-out infinite;"></div>';
+      } else if (item.src) {
+        var img = document.createElement('img');
+        img.src = item.src;
+        img.style.cssText = 'width:100%;height:100%;object-fit:cover;cursor:pointer;display:block;';
+        img.onclick = function() {
+          document.getElementById('numzoo-modal-img').src = item.src;
+          document.getElementById('numzoo-modal').style.display = 'flex';
+        };
+        el.appendChild(img);
+      }
+      wrap.appendChild(el);
+    });
+  }
+
+  document.getElementById('numzoo-modal').onclick = function() { this.style.display='none'; };
+
+  // Poll the collection_trigger component for action payloads
+  var lastTs = -1;
+  setInterval(function() {
+    var el = document.getElementById('numzoo-coll-data');
+    if (!el) return;
+    try {
+      var data = JSON.parse(el.dataset.payload||'{}');
+      if (!data.ts || data.ts === lastTs) return;
+      lastTs = data.ts;
+      (data.actions||[]).forEach(function(act) {
+        if (act.action === 'add-locked') window.numzooAddLocked(act.id);
+        else if (act.action === 'unlock')    window.numzooUnlock(act.id, act.src);
+        else if (act.action === 'failed')    window.numzooMarkFailed(act.id);
+      });
+    } catch(e) {}
+  }, 200);
+
+  setTimeout(renderCollection, 800);
+})();
+</script>
 """
 
 with gr.Blocks(title="🦁 NumZoo") as demo:
 
-    state = gr.State({})
+    state           = gr.State({})
+    hidden_image    = gr.Image(visible=False, label="", type="pil")
+    hidden_data_url = gr.Textbox(visible=False, value="")
 
     gr.Markdown("# 🦁 NumZoo", elem_id="title")
     gr.Markdown("*Do maths. Win cute animals!*", elem_id="subtitle")
@@ -263,22 +459,19 @@ with gr.Blocks(title="🦁 NumZoo") as demo:
     gr.HTML("""
 <script>
   document.addEventListener('DOMContentLoaded', function() {
-    const saved = localStorage.getItem('numzoo_name');
+    var saved = localStorage.getItem('numzoo_name');
     if (saved) {
-      setTimeout(() => {
-        const input = document.querySelector('input[placeholder="Your name…"]');
+      setTimeout(function() {
+        var input = document.querySelector('input[placeholder="Your name…"]');
         if (input) {
-          const nativeInputValueSetter = Object.getOwnPropertyDescriptor(window.HTMLInputElement.prototype, 'value').set;
-          nativeInputValueSetter.call(input, saved);
-          input.dispatchEvent(new Event('input', { bubbles: true }));
+          Object.getOwnPropertyDescriptor(HTMLInputElement.prototype,'value').set.call(input, saved);
+          input.dispatchEvent(new Event('input', {bubbles:true}));
         }
       }, 800);
     }
   });
   document.addEventListener('input', function(e) {
-    if (e.target.placeholder === 'Your name…') {
-      localStorage.setItem('numzoo_name', e.target.value);
-    }
+    if (e.target.placeholder === 'Your name…') localStorage.setItem('numzoo_name', e.target.value);
   });
 </script>
 """)
@@ -313,96 +506,14 @@ with gr.Blocks(title="🦁 NumZoo") as demo:
 
         with gr.Group(visible=False) as reward_panel:
             gr.Markdown("### 🎁 Your reward!")
-            loader_html  = gr.HTML(LOADER_HTML, visible=False)
+            loader_html  = gr.HTML("")                            # empty = hidden; content = shown
             reward_image = gr.Image(label="", show_label=False,
-                                    elem_id="reward-img", height=400)
-            reward_error = gr.Markdown("", visible=True)
+                                    elem_id="reward-img", height=400, visible=False)
+            reward_error = gr.Markdown("")
 
     # ── Collection ─────────────────────────────────────────────────────────
-    gr.HTML("""
-<div id="numzoo-collection-wrap" style="margin-top:16px;">
-  <div id="numzoo-collection-header"
-       style="display:none; font-weight:bold; font-size:1.1em; margin:0 0 8px 4px;">🖼️ Collection</div>
-  <div id="numzoo-collection"
-       style="display:flex; flex-wrap:wrap; gap:8px; justify-content:center;"></div>
-</div>
-
-<!-- Zoom modal -->
-<div id="numzoo-modal"
-     style="display:none; position:fixed; inset:0; background:rgba(0,0,0,0.85);
-            z-index:9999; align-items:center; justify-content:center; cursor:pointer;">
-  <img id="numzoo-modal-img"
-       style="max-width:92vw; max-height:92vh; border-radius:20px; box-shadow:0 8px 40px #0008;">
-</div>
-
-<script>
-(function() {
-  const STORAGE_KEY = 'numzoo_collection';
-  const MAX_IMAGES  = 30;
-
-  function getCollection() {
-    try { return JSON.parse(localStorage.getItem(STORAGE_KEY) || '[]'); }
-    catch { return []; }
-  }
-
-  function renderCollection() {
-    const col    = getCollection();
-    const wrap   = document.getElementById('numzoo-collection');
-    const header = document.getElementById('numzoo-collection-header');
-    if (!wrap) return;
-    wrap.innerHTML = '';
-    if (col.length === 0) { header.style.display = 'none'; return; }
-    header.style.display = 'block';
-    [...col].reverse().forEach(src => {
-      const img = document.createElement('img');
-      img.src   = src;
-      Object.assign(img.style, {
-        width: '80px', height: '80px', objectFit: 'cover',
-        borderRadius: '12px', cursor: 'pointer',
-        border: '2px solid transparent', transition: 'border-color .15s'
-      });
-      img.onmouseenter = () => img.style.borderColor = '#7c3aed';
-      img.onmouseleave = () => img.style.borderColor = 'transparent';
-      img.onclick = () => {
-        document.getElementById('numzoo-modal-img').src = src;
-        document.getElementById('numzoo-modal').style.display = 'flex';
-      };
-      wrap.appendChild(img);
-    });
-  }
-
-  document.getElementById('numzoo-modal').onclick = function() {
-    this.style.display = 'none';
-  };
-
-  // Poll for new reward image and save to localStorage via canvas
-  let lastSaved = '';
-  setInterval(() => {
-    const img = document.querySelector('#reward-img img');
-    if (!img || !img.complete || !img.src || img.src === lastSaved) return;
-    if (img.src.includes('.svg') || img.naturalWidth === 0) return;
-    lastSaved = img.src;
-    try {
-      const canvas = document.createElement('canvas');
-      canvas.width  = img.naturalWidth;
-      canvas.height = img.naturalHeight;
-      canvas.getContext('2d').drawImage(img, 0, 0);
-      const dataUrl = canvas.toDataURL('image/jpeg', 0.75);
-      const col = getCollection();
-      if (!col.includes(dataUrl)) {
-        col.push(dataUrl);
-        localStorage.setItem(STORAGE_KEY, JSON.stringify(col.slice(-MAX_IMAGES)));
-        renderCollection();
-      }
-    } catch(e) { console.warn('NumZoo collection save error:', e); }
-  }, 800);
-
-  // Render on page load
-  document.addEventListener('DOMContentLoaded', renderCollection);
-  setTimeout(renderCollection, 1200);  // fallback after Gradio hydrates
-})();
-</script>
-""")
+    collection_trigger = gr.HTML("")   # receives JSON-encoded collection action payloads
+    gr.HTML(COLLECTION_HTML)
 
     restart_btn = gr.Button("🔄 Restart", variant="secondary", size="sm")
 
@@ -415,25 +526,37 @@ with gr.Blocks(title="🦁 NumZoo") as demo:
                        [state, welcome_panel, emoji_panel, game_panel,
                         animal_picker, place_picker])
 
-    go_btn.click(start_game, [animal_picker, place_picker, state],
-                 [state, emoji_panel, game_panel, picker_error,
-                  status_md, question_md])
-
-    # check_answer → then → generate_image (chained, runs immediately after)
-    check_outputs = [state, status_md, question_md, answer_input,
-                     feedback_md, reward_panel, loader_html, reward_image, reward_error]
-    gen_outputs   = [loader_html, reward_image, reward_error]
-
-    check_btn.click(check_answer, [answer_input, state], check_outputs).then(
-        generate_image, inputs=[state], outputs=gen_outputs
-    )
-    answer_input.submit(check_answer, [answer_input, state], check_outputs).then(
-        generate_image, inputs=[state], outputs=gen_outputs
+    # "Let's go!" → start game → immediately pre-generate first reward image
+    go_btn.click(
+        start_game, [animal_picker, place_picker, state],
+        [state, emoji_panel, game_panel, picker_error, status_md, question_md, collection_trigger]
+    ).then(
+        pregenerate_image, [state],
+        [state, hidden_image, hidden_data_url, collection_trigger]
     )
 
+    check_outputs = [
+        state, status_md, question_md, answer_input,
+        feedback_md, reward_panel, loader_html, reward_image, reward_error,
+        hidden_image, hidden_data_url, collection_trigger,
+    ]
+    ondemand_outputs = [state, loader_html, reward_image, reward_error, collection_trigger]
+    pre_outputs      = [state, hidden_image, hidden_data_url, collection_trigger]
+
+    for trigger in [check_btn.click, answer_input.submit]:
+        trigger(
+            check_answer,
+            [answer_input, state, hidden_image, hidden_data_url],
+            check_outputs,
+        ).then(
+            generate_on_demand, [state], ondemand_outputs
+        ).then(
+            pregenerate_image,  [state], pre_outputs
+        )
+
     restart_btn.click(restart, [state],
                       [state, welcome_panel, emoji_panel, game_panel,
-                       player_name, picker_error])
+                       player_name, picker_error, hidden_image, hidden_data_url])
 
 
 if __name__ == "__main__":