Update app.py

app.py

@@ -1,12 +1,3 @@
-# -*- coding: utf-8 -*-
-"""S-D-XL-pipeline.ipynb
-
-Automatically generated by Colaboratory.
-
-Original file is located at
-    https://colab.research.google.com/drive/1iDFfGVa5XpZSZ1yzv-HeFpiB763xt67D
-"""
-
 import requests
 import gradio as gr
 from PIL import Image
@@ -15,7 +6,6 @@ from transformers import utils
 utils.move_cache()
 import os
 
-
 hf_token = os.getenv('HF_TOKEN')
 
 API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
@@ -23,12 +13,24 @@ headers = {"Authorization": f"Bearer {hf_token}"}
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
+    if response.status_code != 200:
+        # If the response is not successful, return None or handle it accordingly
+        print("Error from the API:", response.text)  # For debugging
+        return None
     return response.content
 
 def generate_image(prompt):
     image_bytes = query({"inputs": prompt})
-    image = Image.open(io.BytesIO(image_bytes))
-    return image
+    if image_bytes is None:
+        # Handle the case where the API did not return image data
+        return "The API call was unsuccessful. Check the logs for details."
+    
+    try:
+        image = Image.open(io.BytesIO(image_bytes))
+        return image
+    except IOError:
+        # Handle cases where PIL cannot open the bytes received
+        return "The returned data could not be recognized as an image."
 
 iface = gr.Interface(
     fn=generate_image,
@@ -38,4 +40,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-