File size: 1,041 Bytes
b0a393a
 
 
 
d2e7948
 
b0a393a
d3c69c2
b0a393a
 
d3c69c2
d2e7948
 
 
b0a393a
 
 
 
 
 
 
 
 
 
 
 
d2e7948
b0a393a
d2e7948
b0a393a
 
d2e7948
b0a393a
d2e7948
b0a393a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import Response
from io import BytesIO
from PIL import Image
import torch
import uvicorn
import os

# Import the CodeFormer model processing function
from codeformer_model import enhance_image  # Make sure this function is defined

app = FastAPI()

@app.post("/enhance")
async def enhance_image_api(file: UploadFile = File(...)):
    try:
        # Load image
        image = Image.open(file.file).convert("RGB")

        # Process the image using the CodeFormer model
        enhanced_image = enhance_image(image)

        # Convert the processed image to bytes
        img_byte_arr = BytesIO()
        enhanced_image.save(img_byte_arr, format="PNG")
        img_byte_arr = img_byte_arr.getvalue()

        return Response(content=img_byte_arr, media_type="image/png")

    except Exception as e:
        return {"error": str(e)}

# Required to run on Hugging Face Spaces
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))