Update app.py

app.py

@@ -4,7 +4,7 @@ import torchvision.transforms as transforms
 from PIL import Image
 import os
 import requests
-from io import BytesIO
+import numpy as np
 
 # Define the model URL and local path
 MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
@@ -22,8 +22,9 @@ if not os.path.exists(MODEL_PATH):
 # Define the Real-ESRGAN class
 class RealESRGAN:
     def __init__(self, model_path, device):
-        self.model = torch.hub.load('xinntao/Real-ESRGAN', 'restoration', model=model_path, device=device)
-        self.model.eval()
+        # Load the model using torch.hub
+        self.model = torch.hub.load('xinntao/Real-ESRGAN', 'real_esrgan', model=model_path, force_reload=True)
+        self.model.eval().to(device)  # Put model in evaluation mode and move to device
         
     def upscale(self, image_tensor):
         with torch.no_grad():
@@ -44,15 +45,15 @@ def upscale_image(image):
     sr_image_tensor = model.upscale(image_tensor)
 
     # Post-process the output image and convert back to PIL
-    sr_image = (sr_image_tensor.squeeze(0).numpy() * 255).astype(np.uint8)
+    sr_image = (sr_image_tensor.squeeze(0).cpu().numpy() * 255).astype(np.uint8)
     sr_image = Image.fromarray(sr_image.transpose(1, 2, 0))  # Change to HWC format
     return sr_image
 
 # Create a Gradio interface
 interface = gr.Interface(
     fn=upscale_image,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs=gr.outputs.Image(type="pil"),
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
     title="Image Upscaler with R-ESRGAN Anime 6B",
     description="Upload an image, and it'll be upscaled using the R-ESRGAN Anime model. Supports inputs in various formats."
 )