Update app.py

app.py

@@ -1,20 +1,23 @@
-import os
-import sys
+'''
+    Gradio demo (almost the same code as the one used in Huggingface space)
+'''
+import os, sys
 import cv2
 import time
-import datetime
-import pytz
+import datetime, pytz
 import gradio as gr
 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)
 from test_code.inference import super_resolve_img
 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
@@ -22,86 +25,62 @@ def auto_download_if_needed(weight_path):
     if not os.path.exists("pretrained"):
         os.makedirs("pretrained")
     
-    weight_mappings = {
-        "pretrained/4x_APISR_RRDB_GAN_generator.pth": "v0.2.0/4x_APISR_RRDB_GAN_generator.pth",
-        "pretrained/4x_APISR_GRL_GAN_generator.pth": "v0.1.0/4x_APISR_GRL_GAN_generator.pth",
-        "pretrained/2x_APISR_RRDB_GAN_generator.pth": "v0.1.0/2x_APISR_RRDB_GAN_generator.pth",
-        "pretrained/4x_APISR_DAT_GAN_generator.pth": "v0.3.0/4x_APISR_DAT_GAN_generator.pth"
-    }
-    
-    if weight_path in weight_mappings:
-        version_path = weight_mappings[weight_path]
-        filename = os.path.basename(weight_path)
-        os.system(f"wget https://github.com/Kiteretsu77/APISR/releases/download/{version_path}")
-        os.system(f"mv {filename} pretrained")
-
-def load_model_with_device(loader_func, weight_path, scale, device):
-    # Check if CUDA is available, otherwise fall back to CPU
-    if device.type == 'cuda' and not torch.cuda.is_available():
-        device = torch.device('cpu')  # Force CPU if CUDA is not available
+    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")
     
-    # Load the state dict with map_location to handle CPU if necessary
-    state_dict = torch.load(weight_path, map_location=device)
-    
-    # Initialize the model
-    generator = loader_func(weight_path, scale=scale)
+    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")
     
-    # Load the state dict and move to device
-    if hasattr(generator, 'load_state_dict'):
-        generator.load_state_dict(state_dict)
-    generator = generator.to(device)
+    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")
     
-    return generator
+
 
 def inference(img_path, model_name):
+    
     try:
-        # Determine device - use GPU if available, otherwise CPU
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         weight_dtype = torch.float32
         
-        # Model configurations
-        model_configs = {
-            "4xGRL": ("pretrained/4x_APISR_GRL_GAN_generator.pth", load_grl, 4),
-            "4xRRDB": ("pretrained/4x_APISR_RRDB_GAN_generator.pth", load_rrdb, 4),
-            "2xRRDB": ("pretrained/2x_APISR_RRDB_GAN_generator.pth", load_rrdb, 2),
-            "4xDAT": ("pretrained/4x_APISR_DAT_GAN_generator.pth", load_dat, 4)
-        }
-        
-        if model_name not in model_configs:
-            raise gr.Error("Unsupported model selected")
+        # 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)  # Directly use default way now
+            
+        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
+            
+        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
+            
+        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
             
-        weight_path, loader_func, scale = model_configs[model_name]
-        auto_download_if_needed(weight_path)
+        else:
+            raise gr.Error("We don't support such Model")
         
-        # Load model and move to appropriate device
-        try:
-            generator = load_model_with_device(loader_func, weight_path, scale, device)
-        except RuntimeError as e:
-            if "out of memory" in str(e):
-                # If we run out of CUDA memory, try loading on CPU instead
-                device = torch.device('cpu')
-                generator = load_model_with_device(loader_func, weight_path, scale, device)
-            else:
-                raise e
-
         generator = generator.to(dtype=weight_dtype)
-        generator.eval()  # Set to evaluation mode
-
-        print(f"Processing {img_path} on {device}")
-        print(f"Current time: {datetime.datetime.now(pytz.timezone('US/Eastern'))}")
-
-        # Process 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 convert output
-        store_name = f"output_{time.time()}.png"
+
+
+        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)
+        store_name = str(time.time()) + ".png"
         save_image(super_resolved_img, store_name)
         outputs = cv2.imread(store_name)
         outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
@@ -109,11 +88,16 @@ def inference(img_path, model_name):
         
         return outputs
     
+    
     except Exception as error:
-        raise gr.Error(f"Error during processing: {str(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)
@@ -134,10 +118,15 @@ 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")
 
@@ -145,16 +134,19 @@ if __name__ == '__main__':
                 output_image = gr.Image(type="numpy", label="Output image")
 
         with gr.Row(elem_classes=["container"]):
-            gr.Examples([
-                ["__assets__/lr_inputs/image-00277.png"],
-                ["__assets__/lr_inputs/image-00542.png"],
-                ["__assets__/lr_inputs/41.png"],
-                ["__assets__/lr_inputs/f91.jpg"],
-                ["__assets__/lr_inputs/image-00440.png"],
-                ["__assets__/lr_inputs/image-00164.jpg"],
-                ["__assets__/lr_inputs/img_eva.jpeg"],
-                ["__assets__/lr_inputs/naruto.jpg"],
-            ], [input_image])
+            gr.Examples(
+                [
+                    ["__assets__/lr_inputs/image-00277.png"],
+                    ["__assets__/lr_inputs/image-00542.png"],
+                    ["__assets__/lr_inputs/41.png"],
+                    ["__assets__/lr_inputs/f91.jpg"],
+                    ["__assets__/lr_inputs/image-00440.png"],
+                    ["__assets__/lr_inputs/image-00164.jpg"],
+                    ["__assets__/lr_inputs/img_eva.jpeg"],
+                    ["__assets__/lr_inputs/naruto.jpg"],
+                ],
+                [input_image],
+            )
 
         run_btn.click(inference, inputs=[input_image, model_name], outputs=[output_image])