Update app.py

app.py

@@ -13,14 +13,20 @@ async def generate_image(prompt):
         # Generate an image based on the prompt
         output = await asyncio.to_thread(model, prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
-        return output
+
         # Check if the model returned images
         if isinstance(output.images, list) and len(output.images) > 0:
             image = output.images[0]
             buffered = BytesIO()
-            image.save(buffered, format="JPEG")
-            image_bytes = buffered.getvalue()
-            return image_bytes
+            try:
+                image.save(buffered, format="JPEG")
+                image_bytes = buffered.getvalue()
+                # Verify the image bytes
+                print(f"Image bytes length: {len(image_bytes)}")
+                return image_bytes
+            except Exception as e:
+                print(f"Error saving image: {e}")
+                return None
         else:
             raise Exception("No images returned by the model.")
     except Exception as e: