Update main.py

main.py

@@ -1,8 +1,8 @@
-from fastapi import FastAPI, File, UploadFile, Form
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import JSONResponse
 from gradio_client import Client, handle_file
-import shutil
 import base64
+
 import os
 
 app = FastAPI()
@@ -14,34 +14,24 @@ client = Client("Makhinur/Bringingoldphotoliveagain", hf_token=HF_TOKEN)
 
 @app.post("/upload/")
 async def upload_image(file: UploadFile = File(...)):
-    if not file:
-        raise HTTPException(status_code=400, detail="No file uploaded")
-
-    # Save the uploaded file to a temporary location
-    temp_file_path = f"temp_{file.filename}"
-    with open(temp_file_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-
     try:
-        # Use Gradio client to process the image
-        result = client.predict(
-            img=handle_file(temp_file_path),
+        # Save the uploaded file temporarily
+        with open(file.filename, "wb") as buffer:
+            buffer.write(await file.read())
+        
+        # Use the Gradio client to process the image
+        result = gradio_client.predict(
+            img=handle_file(file.filename),
             api_name="/predict"
         )
 
-        # Encode the processed image as base64
-        with open(result[0], "rb") as image_file:
-            encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
-
-        # Clean up the temporary file
-        os.remove(temp_file_path)
+        # Read the output image and encode it in base64
+        with open(result[0], "rb") as result_file:
+            encoded_string = base64.b64encode(result_file.read()).decode('utf-8')
 
-        # Return the processed image as a base64 string
-        return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{encoded_image}"})
+        # Create a JSON response with the base64 encoded image
+        return JSONResponse(content={"sketch_image_base64": f"data:image/jpeg;base64,{encoded_string}"})
 
     except Exception as e:
-        # Clean up the temporary file in case of an error
-        if os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
         raise HTTPException(status_code=500, detail=str(e))