Update app.py

app.py

@@ -34,6 +34,21 @@ def pil_to_b64(img_pil: Image.Image) -> str:
     img_pil.save(buf, format="PNG")
     return base64.b64encode(buf.getvalue()).decode("utf-8")
 
+def _b64_to_bytes(s: str) -> bytes:
+    """
+    Robustly decode base64 / base64url (handles '-' '_' and missing padding).
+    """
+    s = (s or "").strip()
+    s = s.replace("-", "+").replace("_", "/")
+    pad = (-len(s)) % 4
+    if pad:
+        s += "=" * pad
+    try:
+        return base64.b64decode(s, validate=False)
+    except Exception:
+        # final fallback using urlsafe decoder
+        return base64.urlsafe_b64decode(s + "=" * ((4 - len(s) % 4) % 4))
+
 def build_prompt(user_text: str, is_first: bool, lock_longshot: bool = True) -> str:
     """Compose continuity-aware prompt text."""
     user_text = (user_text or "").strip()
@@ -215,15 +230,13 @@ def _poll_and_decode(job_id: str, dbg: list[str]):
                     dbg.append("No 'img' field present.")
                     return None, "\n".join(dbg)
 
-                pad = (-len(b64)) % 4
-                if pad:
-                    b64 = b64 + ("=" * pad)
                 try:
-                    img_bytes = base64.b64decode(b64, validate=False)
-                except binascii.Error as e:
-                    dbg.append(f"Base64 decode error: {e}")
+                    img_bytes = _b64_to_bytes(b64)
+                except Exception as e:
+                    dbg.append(f"Base64/urlsafe decode failed: {type(e).__name__}: {e}")
+                    # Try interpret as URL text after decode
                     try:
-                        txt = base64.b64decode(b64 + '==', validate=False).decode("utf-8", "ignore").strip()
+                        txt = _b64_to_bytes(b64).decode("utf-8", "ignore").strip()
                         if txt.startswith("http"):
                             r = requests.get(txt, timeout=60)
                             r.raise_for_status()
@@ -231,7 +244,7 @@ def _poll_and_decode(job_id: str, dbg: list[str]):
                         else:
                             return None, "\n".join(dbg)
                     except Exception as e2:
-                        dbg.append(f"Secondary b64/text parse failed: {type(e2).__name__}: {e2}")
+                        dbg.append(f"Secondary b64→text URL parse failed: {type(e2).__name__}: {e2}")
                         return None, "\n".join(dbg)
 
                 return _decode_bytes_to_image(img_bytes, dbg)
@@ -293,7 +306,8 @@ def generate_next(prompt_text, steps, size, lock, prev_img, change):
     w, h = _parse_size(size)
     prompt = build_prompt(prompt_text, is_first=False, lock_longshot=lock)
     init_img = prev_img if isinstance(prev_img, Image.Image) else None
-    img, debug = horde_generate(prompt, steps=steps, width=w, height=h, init_image=init_img, denoise=float(change))
+    img, debug = horde_generate(prompt, steps=steps, width=w, height=h,
+                                init_image=init_img, denoise=float(change))
     if img is None:
         gr.Warning("Generation failed. See debug log for details.")
     return img, debug
@@ -310,7 +324,6 @@ def _parse_size(s):
 # =========================
 CUSTOM_CSS = """
 .gradio-container { padding: 24px; }
-
 /* Rounded prompt boxes */
 .prompt-box textarea {
   border-radius: 18px !important;
@@ -319,14 +332,12 @@ CUSTOM_CSS = """
   line-height: 1.4;
   padding: 14px 16px;
 }
-
 /* Pill buttons */
 .pill button {
   border-radius: 999px !important;
   padding: 10px 18px;
   font-size: 15px;
 }
-
 /* Rounded image boxes */
 .image-out .wrap, .image-out .svelte-1ipelgc {
   border-radius: 22px !important;
@@ -412,4 +423,4 @@ with gr.Blocks(css=CUSTOM_CSS, title="Image Checkpoints – Stable Horde (txt2im
 
 if __name__ == "__main__":
     demo.queue().launch()
-    
+