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
@@ -10,7 +7,6 @@ import torch
 import numpy as np
 from torchvision.utils import save_image
 
-
 # Import files from the local folder
 root_path = os.path.abspath('.')
 sys.path.append(root_path)
@@ -25,6 +21,7 @@ def auto_download_if_needed(weight_path):
     if not os.path.exists("pretrained"):
         os.makedirs("pretrained")
     
+    # Tải các mô hình vào CPU
     if weight_path == "pretrained/4x_APISR_RRDB_GAN_generator.pth":
         os.system("wget https://github.com/Kiteretsu77/APISR/releases/download/v0.2.0/4x_APISR_RRDB_GAN_generator.pth")
         os.system("mv 4x_APISR_RRDB_GAN_generator.pth pretrained")
@@ -40,45 +37,64 @@ 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 load_grl_cpu(weight_path, scale=4):
+    # Tải mô hình GRL vào CPU
+    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
+    generator = load_grl(scale=scale)  # Khởi tạo mô hình GRL
+    generator.load_state_dict(state_dict)
+    return generator
+
+def load_rrdb_cpu(weight_path, scale=4):
+    # Tải mô hình RRDB vào CPU
+    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
+    generator = load_rrdb(scale=scale)  # Khởi tạo mô hình RRDB
+    generator.load_state_dict(state_dict)
+    return generator
+
+def load_dat_cpu(weight_path, scale=4):
+    # Tải mô hình DAT vào CPU
+    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
+    generator = load_dat(scale=scale)  # Khởi tạo mô hình DAT
+    generator.load_state_dict(state_dict)
+    return generator
 
 
 def inference(img_path, model_name):
-    
     try:
         weight_dtype = torch.float32
         
-        # Load the model
+        # Load the model based on user selection
         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_cpu(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_cpu(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_cpu(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_cpu(weight_path, scale=4)
             
         else:
             raise gr.Error("We don't support such Model")
         
         generator = generator.to(dtype=weight_dtype)
 
-
         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
+        # Run super-resolution and save result
         super_resolved_img = super_resolve_img(generator, img_path, output_path=None, weight_dtype=weight_dtype, downsample_threshold=720, crop_for_4x=True)
         store_name = str(time.time()) + ".png"
         save_image(super_resolved_img, store_name)
@@ -88,12 +104,10 @@ def inference(img_path, model_name):
         
         return outputs
     
-    
     except Exception as error:
         raise gr.Error(f"global exception: {error}")
 
 
-
 if __name__ == '__main__':
     
     MARKDOWN = \