Update main.py

main.py

@@ -1,53 +1,44 @@
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
 from gradio_client import Client, handle_file
-import os
-import tempfile
+import shutil
 import base64
+import os
 
 app = FastAPI()
 
-# Setup CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Adjust as needed
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Initialize the Gradio client with the token
 client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)
 
+
 @app.post("/upload/")
-async def upload_file(file: UploadFile = File(...), version: str = Form(...), scale: int = Form(...)):
-    # Save the uploaded file temporarily
-    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-        temp_file.write(await file.read())
-        temp_file_path = temp_file.name
+async def enhance_image(
+    file: UploadFile = File(...),
+    version: str = Form(...),
+    scale: int = Form(...)
+):
+    # Save the uploaded image to a temporary file
+    temp_file_path = "temp_image.png"
+    with open(temp_file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
 
     try:
-        # Use handle_file to prepare the file for the Gradio client
-        result = client.predict(handle_file(temp_file_path), version, scale, api_name="/predict")
-
-        # Check if the result is valid
-        if result and len(result) == 2:
-            # Convert the image data to a base64 string
-            with open(result[0], "rb") as image_file:
-                image_data = base64.b64encode(image_file.read()).decode("utf-8")
-
-            return JSONResponse({
-                "sketch_image_base64": f"data:image/png;base64,{image_data}",
-                "result_file": result[1]
-            })
-        else:
-            return JSONResponse({"error": "Invalid result from the prediction API."}, status_code=500)
+        # Use the Gradio client to process the image
+        result = client.predict(
+            img=handle_file(temp_file_path),
+            version=version,
+            scale=scale,
+            api_name="/predict"
+        )
+
+        # Read the result image and encode it in base64
+        with open(result[0], "rb") as img_file:
+            b64_string = base64.b64encode(img_file.read()).decode('utf-8')
+
+        return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{b64_string}"})
+
     except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-    finally:
-        # Clean up the temporary file
-        if os.path.exists(temp_file_path):
-            os.unlink(temp_file_path)
+        return JSONResponse(status_code=500, content={"message": str(e)})
+