Update main.py

main.py

@@ -4,6 +4,7 @@ from gradio_client import Client, handle_file
 import shutil
 import base64
 import os
+from PIL import Image  # Import the Pillow library
 
 app = FastAPI()
 
@@ -12,7 +13,6 @@ 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",
@@ -43,17 +43,25 @@ async def enhance_image(
             api_name="/predict"
         )
 
-        # Read the result image and encode it in base64
-        with open(result[0], "rb") as img_file:
+        # Assuming the Gradio app outputs a WebP file
+        result_image_path = result[0]  # This path should be a WebP file
+
+        # Convert the WebP image to PNG using Pillow
+        with Image.open(result_image_path) as img:
+            png_image_path = "output_image.png"
+            img.save(png_image_path, format="PNG")
+
+        # Read the PNG image and encode it in base64
+        with open(png_image_path, "rb") as img_file:
             b64_string = base64.b64encode(img_file.read()).decode('utf-8')
 
-        # Clean up the temporary file
+        # Clean up the temporary files
         os.remove(temp_file_path)
+        os.remove(png_image_path)
 
         return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{b64_string}"})
 
     except Exception as e:
         # Log the error message for debugging
         print(f"Error processing image: {e}")
-        return JSONResponse(status_code=500, content={"message": "Internal Server Error"})
-
+        return JSONResponse(status_code=500, content={"message": "Internal Server Error"})