Makhinur commited on
Commit
b18ac6b
·
verified ·
1 Parent(s): 4a667e3

Update main.py

Browse files

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"})

Files changed (1) hide show
  1. main.py +6 -7
main.py CHANGED
@@ -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 image_file:
47
- image_data = base64.b64encode(image_file.read()).decode("utf-8")
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
+