Add better error handling

app.py

@@ -9,7 +9,6 @@ import gradio as gr
 # Make sure you export your API token before running this script:
 #   export CHUTES_API_TOKEN="your_real_token_here"
 # ------------------------------------------------------------------------------
-
 API_TOKEN = os.getenv("CHUTES_API_TOKEN")
 if API_TOKEN is None:
     raise ValueError("Please set the environment variable CHUTES_API_TOKEN before running.")
@@ -21,7 +20,8 @@ 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.
+    and parameters. Returns a PIL.Image on success. Raises a RuntimeError on failure,
+    including if the API returns a non‐JSON response.
     """
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
@@ -42,32 +42,45 @@ def generate_image(prompt: str,
         json=payload
     )
 
+    # If the status code isn’t 200, raise immediately
     if response.status_code != 200:
-        # you can customize error handling as needed
-        raise RuntimeError(f"API returned {response.status_code}: {response.text}")
+        raise RuntimeError(f"API returned HTTP {response.status_code}:\n{response.text}")
+
+    # Try to parse JSON; if that fails, surface raw text
+    try:
+        data = response.json()
+    except Exception as e:
+        raise RuntimeError(f"Failed to parse JSON (status {response.status_code}):\n{response.text}") from e
 
-    data = response.json()
-    # Typically, the API returns a base64‐encoded image under the key "image".
-    # If yours returns a URL (e.g. under "url"), you can fetch that instead.
+    # The API may return base64 under "image", or a URL under "url"/"image_url"
     b64_img = data.get("image")
     if b64_img:
-        decoded = base64.b64decode(b64_img)
-        img = Image.open(io.BytesIO(decoded)).convert("RGB")
-        return 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)}")
 
-    # Fallback: if the API returns a direct URL
     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")
         return img
 
-    raise RuntimeError("No image field found in API response.")
+    # If neither field is present, raise an error
+    raise RuntimeError(f"No 'image' or 'url' field found in API response JSON:\n{data}")
 
 
 # Build the Gradio interface
 with gr.Blocks(title="HiDream Unlimited") as demo:
-    gr.Markdown("## HiDream Unlimited\nGenerate unlimited AI‐driven images powered by HiDream/Chutes.\n\nEnter your prompt and tweak the parameters, then click **Generate**.")
+    gr.Markdown(
+        "## HiDream Unlimited\n\n"
+        "Generate unlimited AI‐driven images powered by HiDream/Chutes.\n\n"
+        "Enter your prompt and tweak the parameters, then click **Generate**."
+    )
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -98,7 +111,6 @@ with gr.Blocks(title="HiDream Unlimited") as demo:
         with gr.Column(scale=1):
             output_img = gr.Image(label="Generated Image", interactive=False)
 
-    # Wire up the button
     generate_btn.click(
         fn=generate_image,
         inputs=[prompt_in, resolution_in, guidance_in, steps_in, shift_in],