Update app.py

app.py

@@ -2,20 +2,60 @@ from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.responses import Response
 from PIL import Image
 import io
+import torch
+import numpy as np
 
 app = FastAPI()
 
+device = torch.device("cpu")
+torch.set_num_threads(1)
+
+models = {}
+
+def get_model(scale: int):
+    if scale not in models:
+        model = torch.hub.load(
+            "eugenesiow/edsr",
+            "edsr_base",
+            scale=scale,
+            pretrained=True,
+            skip_validation=True,
+        )
+        model = model.to(device).eval()
+        models[scale] = model
+    return models[scale]
+
 @app.post("/upscale")
 async def upscale(
     image: UploadFile = File(...),
     scale: int = Form(2)
 ):
-    img = Image.open(io.BytesIO(await image.read()))
+    try:
+        model = get_model(scale)
+    except Exception as e:
+        return Response(
+            f"Model load failed: {type(e).__name__}: {e}",
+            status_code=500,
+            media_type="text/plain",
+        )
+
+    try:
+        img = Image.open(io.BytesIO(await image.read())).convert("RGB")
+
+        with torch.no_grad():
+            sr = model(img)
 
-    # Simple LANCZOS resize — works on every machine, no deps needed
-    new_size = (img.width * scale, img.height * scale)
-    img = img.resize(new_size, Image.LANCZOS)
+        sr_pil = sr.squeeze(0).clamp(0, 1).cpu().numpy()
+        sr_pil = np.transpose(sr_pil, (1, 2, 0))
+        sr_pil = (sr_pil * 255).astype(np.uint8)
+        sr_pil = Image.fromarray(sr_pil)
 
-    buffer = io.BytesIO()
-    img.save(buffer, format="PNG")
-    return Response(buffer.getvalue(), media_type="image/png")
+        buffer = io.BytesIO()
+        sr_pil.save(buffer, format="PNG")
+        return Response(buffer.getvalue(), media_type="image/png")
+    except Exception as e:
+        return Response(
+            f"Upscale failed: {type(e).__name__}: {e}",
+            status_code=500,
+            media_type="text/plain",
+        )