Update app.py

app.py

@@ -1,6 +1,3 @@
-'''
-    Gradio demo (almost the same code as the one used in Huggingface space)
-'''
 import os, sys
 import cv2
 import time
@@ -40,7 +37,6 @@ def auto_download_if_needed(weight_path):
     if weight_path == "pretrained/4x_APISR_DAT_GAN_generator.pth":
         os.system("wget https://github.com/Kiteretsu77/APISR/releases/download/v0.3.0/4x_APISR_DAT_GAN_generator.pth")
         os.system("mv 4x_APISR_DAT_GAN_generator.pth pretrained")
-    
 
 
 def inference(img_path, model_name):
@@ -52,27 +48,28 @@ def inference(img_path, model_name):
         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)  # Directly use default way now
+            generator = load_grl(weight_path, scale=4)
             
         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)  # Directly use default way now
+            generator = load_rrdb(weight_path, scale=4)
             
         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) # Directly use default way now
+            generator = load_rrdb(weight_path, scale=2)
             
         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) # Directly use default way now
+            generator = load_dat(weight_path, scale=4)
             
         else:
             raise gr.Error("We don't support such Model")
         
-        generator = generator.to(dtype=weight_dtype)
+        # Move the model to the CPU
+        generator = generator.to(device='cpu')
 
 
         print("We are processing ", img_path)
@@ -150,4 +147,4 @@ if __name__ == '__main__':
 
         run_btn.click(inference, inputs=[input_image, model_name], outputs=[output_image])
 
-    block.launch()
+    block.launch()