|
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 |
|
from PIL import Image |
|
|
|
app = FastAPI() |
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
|
|
|
client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN) |
|
|
|
|
|
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(...) |
|
): |
|
|
|
gradio_version = version_map.get(version, "v1.4") |
|
|
|
|
|
temp_file_path = "temp_image.png" |
|
with open(temp_file_path, "wb") as buffer: |
|
shutil.copyfileobj(file.file, buffer) |
|
|
|
try: |
|
|
|
result = client.predict( |
|
img=handle_file(temp_file_path), |
|
version=gradio_version, |
|
scale=scale, |
|
api_name="/predict" |
|
) |
|
|
|
|
|
result_image_path = result[0] |
|
|
|
|
|
with Image.open(result_image_path) as img: |
|
png_image_path = "output_image.png" |
|
img.save(png_image_path, format="PNG") |
|
|
|
|
|
with open(png_image_path, "rb") as img_file: |
|
b64_string = base64.b64encode(img_file.read()).decode('utf-8') |
|
|
|
|
|
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: |
|
|
|
print(f"Error processing image: {e}") |
|
return JSONResponse(status_code=500, content={"message": "Internal Server Error"}) |