Update app.py

app.py

@@ -96,11 +96,23 @@ def query(payload, max_retries=5):
     
 # Function for direct prompt to image generation
 def generate_image_from_prompt(prompt):
+     image_bytes = query({"inputs": prompt})
+
+    if image_bytes is None:
+        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
+        d = ImageDraw.Draw(error_img)
+        d.text((10, 150), "Image Generation Failed", fill=(255, 255, 255))
+        return error_img
+
     try:
-        image_bytes = query({"inputs": prompt})
-        return image_bytes
+        image = Image.open(io.BytesIO(image_bytes))
+        return image
     except Exception as e:
-        return f"An error occurred during image generation: {str(e)}"
+        print(f"Error: {e}")
+        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
+        d = ImageDraw.Draw(error_img)
+        d.text((10, 150), "Invalid Image Data", fill=(255, 255, 255))
+        return error_img
 
 
 # Assuming your 'process_audio' and 'generate_image_from_prompt' functions are defined elsewhere