Update app.py

app.py

@@ -3,7 +3,9 @@ import uuid, os
 from PIL import Image
 
 from rembg import remove
-from realesrgan import RealESRGAN
+from realesrgan import RealESRGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
 from gfpgan import GFPGANer
 from diffusers import StableDiffusionImg2ImgPipeline
 import torch
@@ -20,8 +22,28 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # ===== Load models ONCE =====
 
 # Real-ESRGAN
-upscaler = RealESRGAN(torch.device(DEVICE), scale=4)
-upscaler.load_weights("RealESRGAN_x4.pth")
+# ===== Real-ESRGAN =====
+model = RRDBNet(
+    num_in_ch=3,
+    num_out_ch=3,
+    num_feat=64,
+    num_block=23,
+    num_grow_ch=32,
+    scale=4
+)
+
+upscaler = RealESRGANer(
+    scale=4,
+    model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+    model=model,
+    tile=0,
+    tile_pad=10,
+    pre_pad=0,
+    half=(DEVICE == "cuda"),
+    device=DEVICE
+)
+
+#upscaler.load_weights("RealESRGAN_x4.pth")
 
 # GFPGAN
 gfpgan = GFPGANer(
@@ -55,7 +77,7 @@ async def remove_bg(file: UploadFile = File(...)):
 @app.post("/enhance")
 async def enhance(file: UploadFile = File(...)):
     img = Image.open(file.file).convert("RGB")
-    out = upscaler.predict(img)
+    out,_ = upscaler.predict(img)
 
     fname = f"{uuid.uuid4().hex}.png"
     out.save(os.path.join(OUTPUT_DIR, fname))