Update app.py

app.py

@@ -54,12 +54,16 @@ gfpgan = GFPGANer(
     device=DEVICE
 )
 
+
 # Stable Diffusion Img2Img
+dtype = torch.float16 if DEVICE == "cuda" else torch.float32
+
 sd_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
+    torch_dtype=dtype
 ).to(DEVICE)
 
+
 # ===== Endpoints =====
 
 @app.post("/remove-bg")
@@ -77,7 +81,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.enhance(img)
 
     fname = f"{uuid.uuid4().hex}.png"
     out.save(os.path.join(OUTPUT_DIR, fname))
@@ -88,7 +92,7 @@ async def enhance(file: UploadFile = File(...)):
 async def upscale(file: UploadFile = File(...), scale: int = Form(2)):
     img = Image.open(file.file).convert("RGB")
     upscaler.scale = scale
-    out = upscaler.predict(img)
+    out, _ = upscaler.enhance(img, outscale=scale)
 
     fname = f"{uuid.uuid4().hex}.png"
     out.save(os.path.join(OUTPUT_DIR, fname))