Update app.py

app.py

@@ -3,25 +3,35 @@ import torch
 import torchvision.transforms as transforms
 from PIL import Image
 import os
+import requests
+from io import BytesIO
 
-# Download the required model checkpoint if it doesn't exist
-MODEL_PATH = "weights/RealESRGAN_x4 Anime6B.pth"
-if not os.path.exists(MODEL_PATH):
-    os.system(f"wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4_Anime6B.pth -P weights")
+# Define the model URL and local path
+MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4_Anime6B.pth"
+MODEL_PATH = "RealESRGAN_x4_Anime6B.pth"
 
-# Define the Real-ESRGAN class using the pre-trained model path
+# Download the model if it doesn't exist
+if not os.path.exists(MODEL_PATH):
+    response = requests.get(MODEL_URL)
+    if response.status_code == 200:
+        with open(MODEL_PATH, 'wb') as f:
+            f.write(response.content)
+    else:
+        raise Exception(f"Could not download model: {response.status_code}")
+
+# Define the Real-ESRGAN class
 class RealESRGAN:
-    def __init__(self, device):
-        self.model = torch.hub.load("xinntao/Real-ESRGAN", "restoration", model="RealESRGAN_x4_Anime6B.pth", device=device)
+    def __init__(self, model_path, device):
+        self.model = torch.hub.load('xinntao/Real-ESRGAN', 'restoration', model=model_path, device=device)
         self.model.eval()
-
+        
     def upscale(self, image_tensor):
         with torch.no_grad():
             sr_image_tensor = self.model(image_tensor).clamp(0, 1)
         return sr_image_tensor
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = RealESRGAN(device)
+model = RealESRGAN(MODEL_PATH, device)
 
 # Define the image upscaling function
 def upscale_image(image):