Update main.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from gradio_client import Client, handle_file
import shutil
import base64
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(
file: UploadFile = File(...),
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:
shutil.copyfileobj(file.file, buffer)

try:
# Use the Gradio client to process the image
result = client.predict(
img=handle_file(temp_file_path),
version=gradio_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')

# 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:
# Log the error message for debugging
print(f"Error processing image: {e}")
return JSONResponse(status_code=500, content={"message": "Internal Server Error"})

main.py

@@ -1,8 +1,8 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import JSONResponse
 from gradio_client import Client, handle_file
+import shutil
 import base64
-
 import os
 
 app = FastAPI()
@@ -15,23 +15,26 @@ client = Client("Makhinur/Bringingoldphotoliveagain", hf_token=HF_TOKEN)
 @app.post("/upload/")
 async def upload_image(file: UploadFile = File(...)):
     try:
-        # Save the uploaded file temporarily
-        with open(file.filename, "wb") as buffer:
-            buffer.write(await file.read())
-        
+        temporary file
+    temp_file_path = "temp_image.png"
+    with open(temp_file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
         # Use the Gradio client to process the image
-        result = gradio_client.predict(
-            img=handle_file(file.filename),
+        result = client.predict(
+            img=handle_file(temp_file_path),
             api_name="/predict"
         )
 
-        # 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')
+        # 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')
 
-        # Create a JSON response with the base64 encoded image
-        return JSONResponse(content={"sketch_image_base64": f"data:image/jpeg;base64,{encoded_string}"})
+        # Clean up the temporary file
+        os.remove(temp_file_path)
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        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"})