Update app.py

app.py

@@ -132,7 +132,24 @@ async def mainFunction(html_content:str):
         # Optimize the image to be less than 200KB
         await optimize_image(captured_image_path, output_image_file)
 
+        return output_image_file
+
 # Run the asynchronous main function
 @app.post("/convert-html-to-image/")
 async def convert_html_to_image_endpoint(html_content:HtmlInput):
-    await mainFunction(html_content.message)
+    image_path = await mainFunction(html_content.message)
+    
+    # Check if image was generated
+    if not image_path or not os.path.exists(image_path):
+        raise HTTPException(status_code=500, detail="Image generation failed.")
+
+    # Open the image file for streaming
+    with open(image_path, "rb") as img_file:
+        buffer = io.BytesIO(img_file.read())
+
+    buffer.seek(0)
+
+    # Return the image as a downloadable file
+    return StreamingResponse(buffer, media_type="image/jpeg", headers={
+        "Content-Disposition": f"attachment; filename={os.path.basename(image_path)}"
+    })