Update app.py

app.py

@@ -2,6 +2,9 @@ import time, base64, io, os, requests, traceback
 from typing import Optional
 from PIL import Image, UnidentifiedImageError
 import gradio as gr
+import binascii
+import imageio.v2 as imageio
+import numpy as np
 
 # =========================
 # Stable Horde config
@@ -90,7 +93,7 @@ def horde_txt2img(prompt: str,
         dbg.append("SUBMIT exception:\n" + traceback.format_exc())
         return None, "\n".join(dbg)
 
-    # -------- Poll --------
+        # -------- Poll --------
     start = time.time()
     while True:
         try:
@@ -116,17 +119,51 @@ def horde_txt2img(prompt: str,
                 if not gens:
                     dbg.append("DONE but no generations returned.")
                     return None, "\n".join(dbg)
+
                 b64 = gens[0].get("img")
                 if not b64:
                     dbg.append(f"GEN fields: {list(gens[0].keys())}")
                     return None, "\n".join(dbg)
+
+                # ---- robust decode path ----
+                # 1) fix base64 padding if needed
+                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}")
+                    return None, "\n".join(dbg)
+
+                # 2) log magic header to debug
+                head = img_bytes[:12]
+                dbg.append(f"header bytes: {head.hex(' ')}")
+
+                # 3) try Pillow first
                 try:
-                    img_bytes = base64.b64decode(b64)
                     img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
                     return img, "\n".join(dbg)
-                except UnidentifiedImageError:
-                    dbg.append("PIL could not decode image bytes.")
-                    return None, "\n".join(dbg)
+                except Exception as e:
+                    dbg.append(f"PIL decode failed: {type(e).__name__}: {e}")
+
+                # 4) imageio fallback (handles webp/png edge cases)
+                try:
+                    arr = imageio.imread(io.BytesIO(img_bytes))
+                    if isinstance(arr, np.ndarray):
+                        if arr.ndim == 2:  # grayscale → RGB
+                            arr = np.stack([arr, arr, arr], axis=-1)
+                        elif arr.shape[-1] == 4:  # RGBA → RGB
+                            arr = arr[..., :3]
+                        img = Image.fromarray(arr.astype(np.uint8), mode="RGB")
+                        dbg.append("Decoded via imageio fallback.")
+                        return img, "\n".join(dbg)
+                except Exception as e:
+                    dbg.append(f"imageio decode failed: {type(e).__name__}: {e}")
+
+                dbg.append("PIL could not decode image bytes.")
+                return None, "\n".join(dbg)
 
             if time.time() - start > POLL_TIMEOUT:
                 dbg.append("TIMEOUT waiting for Horde.")