Update main.py

main.py

@@ -1,7 +1,8 @@
 from fastapi import FastAPI, UploadFile, Form, File, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from gradio_client import Client, file
+import aiofiles
 import os
 import shutil
 import base64
@@ -29,8 +30,7 @@ os.makedirs(RESULT_FOLDER, exist_ok=True)
 
 @app.post("/")
 async def hello():
-    return {"Wearon":"wearon model is running"}
-
+    return {"Wearon": "wearon model is running"}
 
 @app.post("/process")
 async def predict(product_image_url: str = Form(...), model_image: UploadFile = File(...)):
@@ -40,18 +40,19 @@ async def predict(product_image_url: str = Form(...), model_image: UploadFile =
         
         # Save the uploaded file to the upload directory
         filename = os.path.join(UPLOAD_FOLDER, model_image.filename)
-        with open(filename, "wb") as buffer:
-            shutil.copyfileobj(model_image.file, buffer)
+        async with aiofiles.open(filename, "wb") as buffer:
+            content = await model_image.read()
+            await buffer.write(content)
 
         base_path = os.getcwd()
         full_filename = os.path.normpath(os.path.join(base_path, filename))
 
-        print("Product image = ", product_image_url)
-        print("Model image = ", full_filename)
+        print("Product image =", product_image_url)
+        print("Model image =", full_filename)
         
         # Perform prediction
         try:
-            result = await client.predict(
+            result = client.predict(
                 dict={"background": file(full_filename), "layers": [], "composite": None},
                 garm_img=file(product_image_url),
                 garment_des="Hello!!",
@@ -78,8 +79,8 @@ async def predict(product_image_url: str = Form(...), model_image: UploadFile =
         os.remove(filename)
 
         # Encode the output image in base64
-        with open(local_output_path, "rb") as image_file:
-            encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
+        async with aiofiles.open(local_output_path, "rb") as image_file:
+            encoded_image = base64.b64encode(await image_file.read()).decode('utf-8')
 
         # Return the output image in JSON format
         return JSONResponse(content={"image": encoded_image}, status_code=200)
@@ -94,4 +95,4 @@ async def uploaded_file(filename: str):
     if os.path.exists(file_path):
         return FileResponse(file_path)
     else:
-        raise HTTPException(status_code=404, detail="File not found")
+        raise HTTPException(status_code=404, detail="File not found")