Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import base64
 import io
 import requests
 from PIL import Image
@@ -20,15 +19,15 @@ def generate_image(prompt: str,
                    shift: int) -> Image.Image:
     """
     Calls the HiDream (Chutes) API to generate an image from the given prompt
-    and parameters. Returns a PIL.Image on success. Raises a RuntimeError on failure,
-    including if the API returns a non‐JSON response.
+    and parameters. Assumes the API responds with raw image bytes (e.g. PNG or JPEG)
+    rather than JSON. Returns a PIL.Image on success. Raises a RuntimeError on failure.
     """
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
         "Content-Type": "application/json"
     }
     payload = {
-        "seed": None,                 # always random
+        "seed": None,
         "shift": shift,
         "prompt": prompt,
         "resolution": resolution,
@@ -44,41 +43,27 @@ def generate_image(prompt: str,
 
     # If the status code isn’t 200, raise immediately
     if response.status_code != 200:
-        raise RuntimeError(f"API returned HTTP {response.status_code}:\n{response.text}")
+        # Try to show any text body if present
+        raw_text = response.text.strip()
+        raise RuntimeError(f"API returned HTTP {response.status_code}:\n{raw_text}")
 
-    # Try to parse JSON; if that fails, surface raw text
+    # At this point we assume response.content is raw image bytes.
     try:
-        data = response.json()
-    except Exception as e:
-        raise RuntimeError(f"Failed to parse JSON (status {response.status_code}):\n{response.text}") from e
-
-    # The API may return base64 under "image", or a URL under "url"/"image_url"
-    b64_img = data.get("image")
-    if b64_img:
-        try:
-            decoded = base64.b64decode(b64_img)
-            img = Image.open(io.BytesIO(decoded)).convert("RGB")
-            return img
-        except Exception as e:
-            raise RuntimeError(f"Error decoding base64 image:\n{str(e)}")
-
-    img_url = data.get("url") or data.get("image_url")
-    if img_url:
-        img_response = requests.get(img_url)
-        if img_response.status_code != 200:
-            raise RuntimeError(f"Failed to fetch image URL ({img_url}): HTTP {img_response.status_code}")
-        img = Image.open(io.BytesIO(img_response.content)).convert("RGB")
+        img = Image.open(io.BytesIO(response.content)).convert("RGB")
         return img
-
-    # If neither field is present, raise an error
-    raise RuntimeError(f"No 'image' or 'url' field found in API response JSON:\n{data}")
+    except Exception as e:
+        # If decoding as an image fails, include a hex preview of the first few bytes
+        prefix = response.content[:16].hex()
+        raise RuntimeError(
+            f"Failed to decode image from response content (first 16 bytes: {prefix}...):\n{str(e)}"
+        )
 
 
 # Build the Gradio interface
 with gr.Blocks(title="HiDream Unlimited") as demo:
     gr.Markdown(
         "## HiDream Unlimited\n\n"
-        "Generate unlimited AI‐driven images powered by HiDream/Chutes.\n\n"
+        "Generate unlimited AI-driven images powered by HiDream.\n\n"
         "Enter your prompt and tweak the parameters, then click **Generate**."
     )