Update app.py

app.py

@@ -22,10 +22,11 @@ if not os.path.exists(MODEL_PATH):
 # Define the Real-ESRGAN class
 class RealESRGAN:
     def __init__(self, model_path, device):
-        # 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
-        
+        # Correctly load the model
+        self.model = torch.hub.load('xinntao/Real-ESRGAN', 'real_esrgan', model_path, device=device)
+        self.model.eval()  # Set the model to eval mode
+        self.device = device  # Store the device
+    
     def upscale(self, image_tensor):
         with torch.no_grad():
             sr_image_tensor = self.model(image_tensor).clamp(0, 1)
@@ -39,7 +40,7 @@ def upscale_image(image):
     # Pre-process the input image
     image = image.convert("RGB")
     transform = transforms.ToTensor()
-    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+    image_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
 
     # Perform upscaling
     sr_image_tensor = model.upscale(image_tensor)