Update app.py

app.py

@@ -1,40 +1,35 @@
-import sys
-import os
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import Response
+from io import BytesIO
+from PIL import Image
 import torch
-import gradio as gr
-from fastapi import FastAPI, UploadFile, File
 import uvicorn
-from PIL import Image
-import io
-
-# Ensure the `basicsr` directory is in the system path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(os.path.join(current_dir, "basicsr"))
+import os
 
-# Import the inference function
-from inference_codeformer import inference  
+# Import the CodeFormer model processing function
+from codeformer_model import enhance_image  # Make sure this function is defined
 
 app = FastAPI()
 
-# Load the CodeFormer model
-model_path = "weights/CodeFormer.pth"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 @app.post("/enhance")
-async def enhance_image(file: UploadFile = File(...), upscale: int = 2, fidelity: float = 0.5):
-    """API Endpoint to enhance images using CodeFormer"""
-    image = Image.open(io.BytesIO(await file.read()))
-    image.save("input.png")
+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()
 
-    output_path = inference(
-        input_path="input.png",
-        upscale=upscale,
-        fidelity=fidelity,
-        model_path=model_path,
-        device=device
-    )
+        return Response(content=img_byte_arr, media_type="image/png")
 
-    return {"enhanced_image": f"https://your-space-name.hf.space/{output_path}"}
+    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=7860)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))