Update app.py

app.py

@@ -17,46 +17,38 @@ from test_code.test_utils import load_grl, load_rrdb, load_dat
 def auto_download_if_needed(weight_path):
     if os.path.exists(weight_path):
         return
-    
+
     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")
-    
-    if weight_path == "pretrained/4x_APISR_GRL_GAN_generator.pth":
-        os.system("wget https://github.com/Kiteretsu77/APISR/releases/download/v0.1.0/4x_APISR_GRL_GAN_generator.pth")
-        os.system("mv 4x_APISR_GRL_GAN_generator.pth pretrained")
-        
-    if weight_path == "pretrained/2x_APISR_RRDB_GAN_generator.pth":
-        os.system("wget https://github.com/Kiteretsu77/APISR/releases/download/v0.1.0/2x_APISR_RRDB_GAN_generator.pth")
-        os.system("mv 2x_APISR_RRDB_GAN_generator.pth pretrained")
-    
-    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")
+
+    # Download pretrained weights based on the model type
+    model_weights = {
+        "pretrained/4x_APISR_RRDB_GAN_generator.pth": "https://github.com/Kiteretsu77/APISR/releases/download/v0.2.0/4x_APISR_RRDB_GAN_generator.pth",
+        "pretrained/4x_APISR_GRL_GAN_generator.pth": "https://github.com/Kiteretsu77/APISR/releases/download/v0.1.0/4x_APISR_GRL_GAN_generator.pth",
+        "pretrained/2x_APISR_RRDB_GAN_generator.pth": "https://github.com/Kiteretsu77/APISR/releases/download/v0.1.0/2x_APISR_RRDB_GAN_generator.pth",
+        "pretrained/4x_APISR_DAT_GAN_generator.pth": "https://github.com/Kiteretsu77/APISR/releases/download/v0.3.0/4x_APISR_DAT_GAN_generator.pth"
+    }
+
+    if weight_path in model_weights:
+        os.system(f"wget {model_weights[weight_path]} -P pretrained")
 
 
+# Define functions to load models into CPU if no GPU is available
 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_grl(generator_weight_PATH=weight_path, scale=scale)
     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_rrdb(generator_weight_PATH=weight_path, scale=scale)
     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_dat(generator_weight_PATH=weight_path, scale=scale)
     generator.load_state_dict(state_dict)
     return generator
 
@@ -64,38 +56,43 @@ def load_dat_cpu(weight_path, scale=4):
 def inference(img_path, model_name):
     try:
         weight_dtype = torch.float32
-        
-        # Load the model based on user selection
+
+        # Load the model based on the selected model_name
         if model_name == "4xGRL":
             weight_path = "pretrained/4x_APISR_GRL_GAN_generator.pth"
             auto_download_if_needed(weight_path)
             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_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_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_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')))
 
-        # 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)
+        # Super-resolve the image
+        super_resolved_img = super_resolve_img(
+            generator, img_path, output_path=None,
+            weight_dtype=weight_dtype, downsample_threshold=720, crop_for_4x=True
+        )
+
+        # Save and display the output
         store_name = str(time.time()) + ".png"
         save_image(super_resolved_img, store_name)
         outputs = cv2.imread(store_name)
@@ -103,25 +100,18 @@ def inference(img_path, model_name):
         os.remove(store_name)
         
         return outputs
-    
+
     except Exception as error:
         raise gr.Error(f"global exception: {error}")
 
 
 if __name__ == '__main__':
-    
-    MARKDOWN = \
-    """
+    MARKDOWN = """
     ## <p style='text-align: center'> APISR: Anime Production Inspired Real-World Anime Super-Resolution (CVPR 2024) </p>
-    
     [GitHub](https://github.com/Kiteretsu77/APISR) | [Paper](https://arxiv.org/abs/2403.01598)
-
     APISR aims at restoring and enhancing low-quality low-resolution **anime** images and video sources with various degradations from real-world scenarios.
-    
-    ### Note: Due to memory restriction, all images whose short side is over 720 pixel will be downsampled to 720 pixel with the same aspect ratio.  E.g., 1920x1080 -> 1280x720
-    ### Note: Please check [Model Zoo](https://github.com/Kiteretsu77/APISR/blob/main/docs/model_zoo.md) for the description of each weight and [Here](https://imgsli.com/MjU0MjI0) for model comparisons.
-    
-    ### If APISR is helpful, please help star the [GitHub Repo](https://github.com/Kiteretsu77/APISR). Thanks! ###
+    ### Note: Due to memory restriction, all images whose short side is over 720 pixel will be downsampled to 720 pixel with the same aspect ratio.
+    ### If APISR is helpful, please help star the [GitHub Repo](https://github.com/Kiteretsu77/APISR). Thanks!
     """
 
     block = gr.Blocks().queue(max_size=10)
@@ -132,15 +122,10 @@ if __name__ == '__main__':
             with gr.Column(scale=2):
                 input_image = gr.Image(type="filepath", label="Input")
                 model_name = gr.Dropdown(
-                    [
-                        "2xRRDB",
-                        "4xRRDB",
-                        "4xGRL",
-                        "4xDAT",
-                    ],
+                    ["2xRRDB", "4xRRDB", "4xGRL", "4xDAT"],
                     type="value",
                     value="4xGRL",
-                    label="model",
+                    label="model"
                 )
                 run_btn = gr.Button(value="Submit")