Update main.py
Browse filesfrom 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"})
@@ -12,6 +12,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
|
|
12 |
# Initialize the Gradio client with the token
|
13 |
client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)
|
14 |
|
|
|
15 |
# Version mapping from HTML to Gradio API
|
16 |
version_map = {
|
17 |
"M1": "v1.2",
|
@@ -43,18 +44,16 @@ async def enhance_image(
|
|
43 |
)
|
44 |
|
45 |
# Read the result image and encode it in base64
|
46 |
-
with open(result[0], "rb") as
|
47 |
-
|
48 |
|
49 |
# Clean up the temporary file
|
50 |
os.remove(temp_file_path)
|
51 |
|
52 |
-
return JSONResponse(content={
|
53 |
-
"sketch_image_base64": f"data:image/png;base64,{image_data}",
|
54 |
-
"result_file": result[1]
|
55 |
-
})
|
56 |
|
57 |
except Exception as e:
|
58 |
# Log the error message for debugging
|
59 |
print(f"Error processing image: {e}")
|
60 |
-
return JSONResponse(status_code=500, content={"message": "Internal Server Error"})
|
|
|
|
12 |
# Initialize the Gradio client with the token
|
13 |
client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)
|
14 |
|
15 |
+
|
16 |
# Version mapping from HTML to Gradio API
|
17 |
version_map = {
|
18 |
"M1": "v1.2",
|
|
|
44 |
)
|
45 |
|
46 |
# Read the result image and encode it in base64
|
47 |
+
with open(result[0], "rb") as img_file:
|
48 |
+
b64_string = base64.b64encode(img_file.read()).decode('utf-8')
|
49 |
|
50 |
# Clean up the temporary file
|
51 |
os.remove(temp_file_path)
|
52 |
|
53 |
+
return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{b64_string}"})
|
|
|
|
|
|
|
54 |
|
55 |
except Exception as e:
|
56 |
# Log the error message for debugging
|
57 |
print(f"Error processing image: {e}")
|
58 |
+
return JSONResponse(status_code=500, content={"message": "Internal Server Error"})
|
59 |
+
|