Update app.py

app.py

@@ -1,4 +1,4 @@
-import time, base64, io, os, requests, traceback, binascii
+import time, base64, io, os, requests, traceback, binascii, gzip, zlib, bz2, lzma, re
 from typing import Optional
 from PIL import Image
 import gradio as gr
@@ -13,7 +13,7 @@ HORDE_STATUS = "https://stablehorde.net/api/v2/generate/status/{id}"
 
 # HF Space secret recommended for priority
 HORDE_API_KEY = os.getenv("HORDE_API_KEY", "")
-CLIENT_AGENT = "StitchMaster/0.2 (https://huggingface.co/spaces/your-space)"
+CLIENT_AGENT = "StitchMaster/0.3 (https://huggingface.co/spaces/your-space)"
 
 DEFAULT_STEPS = 24
 DEFAULT_W = 704   # keep defaults under the 715px threshold
@@ -22,6 +22,7 @@ POLL_INTERVAL = 2.5
 POLL_TIMEOUT = 240            # bump if queues are long
 MODEL = None                  # or set e.g. "SDXL 1.0"
 
+# ---------------- Helpers ----------------
 def _headers():
     # Always send an apikey; fallback to anonymous for testing
     return {
@@ -46,9 +47,35 @@ def _b64_to_bytes(s: str) -> bytes:
     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 _maybe_decompress(buf: bytes, dbg: list[str]) -> bytes:
+    """
+    If bytes are compressed, try gzip, zlib, bz2, lzma (in that order).
+    Returns original buf if none succeed.
+    """
+    head = buf[:4]
+    try:
+        # gzip
+        if len(buf) > 2 and buf[0:2] == b"\x1f\x8b":
+            dbg.append("Detected gzip; decompressing…")
+            return gzip.decompress(buf)
+        # zlib (78 01/9C/DA)
+        if len(buf) > 2 and buf[0] == 0x78 and buf[1] in (0x01, 0x5E, 0x9C, 0xDA):
+            dbg.append("Detected zlib; decompressing…")
+            return zlib.decompress(buf)
+        # bz2
+        if len(buf) > 3 and buf[0:3] == b"BZh":
+            dbg.append("Detected bz2; decompressing…")
+            return bz2.decompress(buf)
+        # lzma/xz
+        if len(buf) > 6 and buf[0:6] == b"\xfd7zXZ\x00":
+            dbg.append("Detected lzma/xz; decompressing…")
+            return lzma.decompress(buf)
+    except Exception as e:
+        dbg.append(f"Decompress probe failed: {type(e).__name__}: {e}")
+    return buf
+
 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()
@@ -72,7 +99,7 @@ def build_prompt(user_text: str, is_first: bool, lock_longshot: bool = True) ->
     return base
 
 # =========================
-# Horde client with debugging (txt2img OR img2img)
+# Horde client (txt2img OR img2img)
 # =========================
 def horde_generate(
     prompt: str,
@@ -105,7 +132,7 @@ def horde_generate(
                 payload["params"]["steps"]  = min(int(payload["params"]["steps"]), 30)
                 payload["params"]["width"]  = min(int(payload["params"]["width"]), 704)
                 payload["params"]["height"] = min(int(payload["params"]["height"]), 704)
-                dbg.append("Fallback applied: steps<=30, width/height<=704. Retrying submit...")
+                dbg.append("Fallback applied: steps<=30, width/height<=704. Retrying submit…")
                 sub = requests.post(HORDE_URL, json=payload, headers=_headers(), timeout=30)
         return sub
 
@@ -143,7 +170,6 @@ def horde_generate(
                 dbg.append(f"SUBMIT json={submit_j}")
                 raise gr.Error("Horde submit succeeded but no job id returned.")
             dbg.append(f"JOB id={job_id}")
-            # Poll & decode
             return _poll_and_decode(job_id, dbg)
         except Exception:
             dbg.append("IMG2IMG path failed, falling back to text-only:\n" + traceback.format_exc())
@@ -211,8 +237,14 @@ def _poll_and_decode(job_id: str, dbg: list[str]):
                 dbg.append(f"GEN keys: {list(g0.keys())}")
                 dbg.append(f"img_type: {g0.get('img_type')}")
 
-                # Prefer URL if present
-                url = g0.get("r2") or g0.get("url") or g0.get("src") or g0.get("image_url")
+                # Prefer URL if present (also check gen_metadata for r2 url)
+                url = (
+                    g0.get("r2")
+                    or g0.get("url")
+                    or g0.get("src")
+                    or g0.get("image_url")
+                    or (isinstance(g0.get("gen_metadata"), dict) and (g0["gen_metadata"].get("r2") or g0["gen_metadata"].get("url")))
+                )
                 if isinstance(url, str) and (url.startswith("http://") or url.startswith("https://")):
                     dbg.append("Found URL in generation → fetching…")
                     try:
@@ -224,28 +256,46 @@ def _poll_and_decode(job_id: str, dbg: list[str]):
                         dbg.append(f"URL fetch failed: {type(e).__name__}: {e}")
                         return None, "\n".join(dbg)
 
-                # Base64 branch
+                # Base64 (or hex/encoded) path
                 b64 = g0.get("img")
                 if not b64:
                     dbg.append("No 'img' field present.")
                     return None, "\n".join(dbg)
 
+                # If 'img' looks like a URL string, 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
+                        return _decode_bytes_to_image(img_bytes, dbg)
+                    except Exception as e:
+                        dbg.append(f"URL fetch failed: {type(e).__name__}: {e}")
+                        return None, "\n".join(dbg)
+
+                # Try base64-url safe first
                 try:
                     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 = _b64_to_bytes(b64).decode("utf-8", "ignore").strip()
-                        if txt.startswith("http"):
-                            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 URL parse failed: {type(e2).__name__}: {e2}")
-                        return None, "\n".join(dbg)
+                    img_bytes = None
+
+                # If that failed or header looks wrong, try hex
+                if not img_bytes or len(img_bytes) < 8:
+                    if re.fullmatch(r"[0-9a-fA-F]+", b64) and len(b64) % 2 == 0:
+                        dbg.append("img looks like hex → decoding…")
+                        try:
+                            img_bytes = bytes.fromhex(b64)
+                        except Exception as e:
+                            dbg.append(f"Hex decode failed: {type(e).__name__}: {e}")
+                            img_bytes = None
+
+                if not img_bytes:
+                    return None, "\n".join(dbg)
+
+                # Some workers compress payloads—try to decompress if needed
+                img_bytes = _maybe_decompress(img_bytes, dbg)
 
                 return _decode_bytes_to_image(img_bytes, dbg)
 
@@ -422,5 +472,4 @@ with gr.Blocks(css=CUSTOM_CSS, title="Image Checkpoints – Stable Horde (txt2im
     b4.click(fn=generate_next,     inputs=[p4, steps, size, lock, img3, change], outputs=[img4, debug_box])
 
 if __name__ == "__main__":
-    demo.queue().launch()
-  
+    demo.queue().launch()