upscale / app.py
darkvibe314's picture
Update app.py
1a08837 verified
Raw
History Blame Contribute Delete
1.87 kB
import io
import cv2
import numpy as np
import requests
import os
import gc
import torch
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response
from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet
app = FastAPI()
# 🧠 Setup Real-ESRGAN Model
model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
model_path = 'RealESRGAN_x4plus.pth'
if not os.path.exists(model_path):
print("Downloading AI model... please wait.")
response = requests.get(model_url)
with open(model_path, 'wb') as f:
f.write(response.content)
# Initialize the AI Engine once
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
upsampler = RealESRGANer(
scale=4,
model_path=model_path,
model=model,
tile=400, # Crucial for free CPU tier
tile_pad=10,
pre_pad=0,
half=False
)
@app.get("/")
def home():
return {"status": "Silent Neural HD Engine Online"}
@app.post("/upscale")
async def upscale(file: UploadFile = File(...)):
try:
data = await file.read()
nparr = np.frombuffer(data, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
# Process image
output, _ = upsampler.enhance(img, outscale=4)
# Convert to JPG
_, encoded_img = cv2.imencode('.jpg', output, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
# Cleanup memory immediately
del img
del output
gc.collect()
return Response(content=encoded_img.tobytes(), media_type="image/jpeg")
except Exception as e:
print(f"Error: {e}")
return {"error": "Processing failed. Image might be too large."}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)