Update app.py

app.py

@@ -19,35 +19,46 @@ def get_model(scale: int):
     return models[scale]
 
 
+def _upscale_image(img: Image.Image, scale: int) -> Image.Image:
+    model = get_model(scale)
+    inputs = ImageLoader.load_image(img)
+    pred = model(inputs)
+    fd, path = tempfile.mkstemp(suffix=".png")
+    try:
+        ImageLoader.save_image(pred, path)
+        os.close(fd)
+        result = Image.open(path).convert("RGB")
+        result.load()
+    finally:
+        os.unlink(path)
+    return result
+
+
 @app.post("/upscale")
 async def upscale(image: UploadFile = File(...), scale: int = Form(2)):
     try:
-        model = get_model(scale)
+        img = Image.open(io.BytesIO(await image.read())).convert("RGB")
     except Exception as e:
         return Response(
-            f"Model load failed: {type(e).__name__}: {e}",
+            f"Image load failed: {type(e).__name__}: {e}",
             status_code=500,
             media_type="text/plain",
         )
 
     try:
-        img = Image.open(io.BytesIO(await image.read())).convert("RGB")
-        inputs = ImageLoader.load_image(img)
-        pred = model(inputs)
-
-        fd, path = tempfile.mkstemp(suffix=".png")
-        try:
-            ImageLoader.save_image(pred, path)
-            os.close(fd)
-            with open(path, "rb") as f:
-                data = f.read()
-        finally:
-            os.unlink(path)
-
-        return Response(data, media_type="image/png")
+        if scale == 8:
+            # Chain: 4x → 2x
+            img = _upscale_image(img, 4)
+            img = _upscale_image(img, 2)
+        else:
+            img = _upscale_image(img, scale)
     except Exception as e:
         return Response(
             f"Upscale failed: {type(e).__name__}: {e}",
             status_code=500,
             media_type="text/plain",
         )
+
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return Response(buf.getvalue(), media_type="image/png")