Update app.py

app.py

@@ -42,41 +42,36 @@ def auto_download_if_needed(weight_path):
 def inference(img_path, model_name):
     
     try:
-        weight_dtype = torch.float32
-        
         # Load the model
         if model_name == "4xGRL":
             weight_path = "pretrained/4x_APISR_GRL_GAN_generator.pth"
             auto_download_if_needed(weight_path)
-            generator = load_grl(weight_path, scale=4)
+            generator = load_grl(weight_path, scale=4, device='cpu')
             
         elif model_name == "4xRRDB":
             weight_path = "pretrained/4x_APISR_RRDB_GAN_generator.pth"
             auto_download_if_needed(weight_path)
-            generator = load_rrdb(weight_path, scale=4)
+            generator = load_rrdb(weight_path, scale=4, device='cpu')
             
         elif model_name == "2xRRDB":
             weight_path = "pretrained/2x_APISR_RRDB_GAN_generator.pth"
             auto_download_if_needed(weight_path)
-            generator = load_rrdb(weight_path, scale=2)
+            generator = load_rrdb(weight_path, scale=2, device='cpu')
             
         elif model_name == "4xDAT":
             weight_path = "pretrained/4x_APISR_DAT_GAN_generator.pth"
             auto_download_if_needed(weight_path)
-            generator = load_dat(weight_path, scale=4)
+            generator = load_dat(weight_path, scale=4, device='cpu')
             
         else:
             raise gr.Error("We don't support such Model")
-        
-        # Move the model to the CPU
-        generator = generator.to(device='cpu')
 
 
         print("We are processing ", img_path)
         print("The time now is ", datetime.datetime.now(pytz.timezone('US/Eastern')))
 
         # In default, we will automatically use crop to match 4x size
-        super_resolved_img = super_resolve_img(generator, img_path, output_path=None, weight_dtype=weight_dtype, downsample_threshold=720, crop_for_4x=True)
+        super_resolved_img = super_resolve_img(generator, img_path, output_path=None, downsample_threshold=720, crop_for_4x=True)
         store_name = str(time.time()) + ".png"
         save_image(super_resolved_img, store_name)
         outputs = cv2.imread(store_name)