Update main.py

main.py

@@ -7,11 +7,18 @@ import os
 
 app = FastAPI()
 
+
 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)
 
+# Version mapping from HTML to Gradio API
+version_map = {
+    "M1": "v1.2",
+    "M2": "v1.3",
+    "M3": "v1.4"
+}
 
 @app.post("/upload/")
 async def enhance_image(
@@ -19,6 +26,9 @@ async def enhance_image(
     version: str = Form(...),
     scale: int = Form(...)
 ):
+    # Map version from HTML to Gradio expected value
+    gradio_version = version_map.get(version, "v1.4")
+
     # Save the uploaded image to a temporary file
     temp_file_path = "temp_image.png"
     with open(temp_file_path, "wb") as buffer:
@@ -28,7 +38,7 @@ async def enhance_image(
         # Use the Gradio client to process the image
         result = client.predict(
             img=handle_file(temp_file_path),
-            version=version,
+            version=gradio_version,
             scale=scale,
             api_name="/predict"
         )
@@ -37,8 +47,13 @@ async def enhance_image(
         with open(result[0], "rb") as img_file:
             b64_string = base64.b64encode(img_file.read()).decode('utf-8')
 
+        # Clean up the temporary file
+        os.remove(temp_file_path)
+
         return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{b64_string}"})
 
     except Exception as e:
-        return JSONResponse(status_code=500, content={"message": str(e)})
+        # Log the error message for debugging
+        print(f"Error processing image: {e}")
+        return JSONResponse(status_code=500, content={"message": "Internal Server Error"})