Spaces:
Sleeping
Sleeping
File size: 1,717 Bytes
8e8e507 024be93 2a74393 21e180e 2a74393 ff0dc43 8e8e507 ff0dc43 21e180e 3727d7a ff0dc43 8e8e507 006e3cd 2a74393 8e8e507 ff0dc43 006e3cd ff0dc43 006e3cd ff0dc43 006e3cd ff0dc43 8e8e507 006e3cd | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | import os
import tempfile
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import Response
from PIL import Image
import io
from super_image import MsrnModel, ImageLoader
app = FastAPI()
models = {}
def get_model(scale: int):
if scale not in models:
model = MsrnModel.from_pretrained("eugenesiow/msrn", scale=scale)
model = model.eval()
models[scale] = model
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:
img = Image.open(io.BytesIO(await image.read())).convert("RGB")
except Exception as e:
return Response(
f"Image load failed: {type(e).__name__}: {e}",
status_code=500,
media_type="text/plain",
)
try:
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") |