Update app.py

app.py

@@ -1,8 +1,7 @@
-import time, base64, io, os, requests, traceback
+import time, base64, io, os, requests, traceback, binascii
 from typing import Optional
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 import gradio as gr
-import binascii
 import imageio.v2 as imageio
 import numpy as np
 
@@ -52,7 +51,9 @@ def horde_txt2img(prompt: str,
             "n": 1
         },
         "nsfw": False,
-        "censor_nsfw": True
+        "censor_nsfw": True,
+        # Ask Horde to return a CDN URL if available (many deployments support this)
+        "r2": True
     }
     if model:
         payload["models"] = [model]
@@ -93,7 +94,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:
@@ -120,50 +121,65 @@ def horde_txt2img(prompt: str,
                     dbg.append("DONE but no generations returned.")
                     return None, "\n".join(dbg)
 
-                b64 = gens[0].get("img")
+                g0 = gens[0]
+                dbg.append(f"GEN keys: {list(g0.keys())}")
+                dbg.append(f"img_type: {g0.get('img_type')}")
+
+                # Prefer URL fields if present
+                url = g0.get("r2") or g0.get("url") or g0.get("src") or g0.get("image_url")
+                if isinstance(url, str) and (url.startswith("http://") or url.startswith("https://")):
+                    dbg.append("Found URL in generation → fetching…")
+                    try:
+                        r = requests.get(url, timeout=60)
+                        r.raise_for_status()
+                        img_bytes = r.content
+                    except Exception as e:
+                        dbg.append(f"URL fetch failed: {type(e).__name__}: {e}")
+                        return None, "\n".join(dbg)
+                    return _decode_bytes_to_image(img_bytes, dbg)
+
+                # Else fall back to base64 field
+                b64 = g0.get("img")
                 if not b64:
-                    dbg.append(f"GEN fields: {list(gens[0].keys())}")
+                    dbg.append("No 'img' field present.")
                     return None, "\n".join(dbg)
 
-                # ---- robust decode path ----
+                # If 'img' looks like URL text (rare), just fetch it
+                if b64.startswith("http://") or b64.startswith("https://"):
+                    dbg.append("img field is a URL string → fetching…")
+                    try:
+                        r = requests.get(b64, timeout=60)
+                        r.raise_for_status()
+                        img_bytes = r.content
+                    except Exception as e:
+                        dbg.append(f"URL fetch failed: {type(e).__name__}: {e}")
+                        return None, "\n".join(dbg)
+                    return _decode_bytes_to_image(img_bytes, dbg)
+
+                # Base64 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 = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-                    return img, "\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)
+                    # Try to interpret as text (maybe it's a URL encoded in base64)
+                    try:
+                        txt = base64.b64decode(b64 + "==", validate=False).decode("utf-8", "ignore").strip()
+                        if txt.startswith("http"):
+                            dbg.append("Base64 decoded to text URL → fetching…")
+                            r = requests.get(txt, timeout=60)
+                            r.raise_for_status()
+                            img_bytes = r.content
+                        else:
+                            return None, "\n".join(dbg)
+                    except Exception as e2:
+                        dbg.append(f"Secondary b64/text parse failed: {type(e2).__name__}: {e2}")
+                        return None, "\n".join(dbg)
+
+                return _decode_bytes_to_image(img_bytes, dbg)
 
             if time.time() - start > POLL_TIMEOUT:
                 dbg.append("TIMEOUT waiting for Horde.")
@@ -175,6 +191,43 @@ def horde_txt2img(prompt: str,
             dbg.append("POLL exception:\n" + traceback.format_exc())
             return None, "\n".join(dbg)
 
+def _decode_bytes_to_image(img_bytes: bytes, dbg: list[str]):
+    # Log header
+    head = img_bytes[:12]
+    dbg.append(f"header bytes: {head.hex(' ')}")
+
+    # Try Pillow first
+    try:
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        return img, "\n".join(dbg)
+    except Exception as e:
+        dbg.append(f"PIL decode failed: {type(e).__name__}: {e}")
+
+    # Fallback: imageio
+    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}")
+
+    # Last resort: save bytes for inspection
+    try:
+        tmp = f"unknown_img_{int(time.time())}.bin"
+        with open(tmp, "wb") as f:
+            f.write(img_bytes)
+        dbg.append(f"Wrote undecodable bytes to {tmp}")
+    except Exception as e:
+        dbg.append(f"Could not write debug bytes: {type(e).__name__}: {e}")
+
+    return None, "\n".join(dbg)
+
 def generate_image(prompt, steps, size):
     # size like "704x704"
     try:
@@ -229,7 +282,7 @@ with gr.Blocks(css=CUSTOM_CSS, title="Image Checkpoints – Stable Horde") as de
     debug_box = gr.Code(label="Debug log", interactive=False)
 
     prompt_boxes, gen_buttons, img_outputs = [], [], []
-    for i in range(1, 4 + 1):
+    for i in range(1, 5):
         with gr.Row():
             with gr.Column(scale=1, min_width=320):
                 p = gr.Textbox(