Update main.py

main.py

@@ -21,6 +21,9 @@ async def load_model():
     global model
     model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
 
+from fastapi.responses import StreamingResponse
+from io import BytesIO
+
 @app.post("/upload/")
 async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
     if model is None:
@@ -34,12 +37,11 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
     aligned_face, instyle, message = model.detect_and_align_image("uploaded_image.jpg", top, bottom, left, right)
     processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
 
-    # Save the processed image
-    with open("result_image.jpg", "wb") as result_buffer:
-        result_buffer.write(processed_image)
+    # Convert processed image to bytes
+    image_bytes = cv2.imencode('.jpg', processed_image)[1].tobytes()
 
-    # Return the processed image
-    return FileResponse("result_image.jpg", media_type="image/jpeg", headers={"Content-Disposition": "attachment; filename=result_image.jpg"})
+    # Return the processed image as a streaming response
+    return StreamingResponse(BytesIO(image_bytes), media_type="image/jpeg")
 
 
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")