Update main.py

main.py

@@ -1,49 +1,52 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
+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 base64
 
-# Initialize FastAPI
 app = FastAPI()
 
-# Load environment variables if needed
+# 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/Its", hf_token=HF_TOKEN)
+client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)
 
 @app.post("/upload/")
-async def upload_image(file: UploadFile = File(...)):
-    try:
-        # Save the uploaded file to a temporary location
-        file_location = f"temp_{file.filename}"
-        with open(file_location, "wb") as f:
-            f.write(await file.read())
-
-        # Use the Gradio client to send the image to the Gradio app
-        result = client.predict(
-            img=handle_file(file_location),
-            api_name="/predict"
-        )
-
-        # Clean up the temporary file
-        os.remove(file_location)
-
-        # Read the sketch image from the result
-        sketch_image_path = result[0]  # Path to the sketch image
-
-        # Convert the sketch image to base64
-        with open(sketch_image_path, "rb") as img_file:
-            sketch_image_base64 = base64.b64encode(img_file.read()).decode('utf-8')
-
-        # Prepare the response
-        response = {
-            "sketch_image_base64": f"data:image/jpeg;base64,{sketch_image_base64}"
-        }
-
-        return JSONResponse(content=response)
+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
 
+    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)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(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)