Update app.py

app.py

@@ -13,105 +13,91 @@ 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
-
+    
     if not os.path.exists("pretrained"):
         os.makedirs("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"
+    
+    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 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):
-    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
-    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):
-    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
-    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):
-    state_dict = torch.load(weight_path, map_location=torch.device('cpu'))
-    generator = load_dat(generator_weight_PATH=weight_path, scale=scale)
-    generator.load_state_dict(state_dict)
-    return generator
-
+    
+    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 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
+        
+        # Load the model with appropriate device mapping
+        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")
+            
+        weight_path, loader_func, scale = model_configs[model_name]
+        auto_download_if_needed(weight_path)
+        
+        # Load model with explicit device mapping
+        generator = loader_func(
+            weight_path, 
+            scale=scale,
+            map_location=device
+        )
+        generator = generator.to(device=device, dtype=weight_dtype)
 
-        # 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')))
+        print(f"Processing {img_path} on {device}")
+        print(f"Current time: {datetime.datetime.now(pytz.timezone('US/Eastern'))}")
 
-        # Super-resolve the image
+        # 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
+            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 and convert output
+        store_name = f"output_{time.time()}.png"
         save_image(super_resolved_img, store_name)
         outputs = cv2.imread(store_name)
         outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
         os.remove(store_name)
         
         return outputs
-
+    
     except Exception as error:
-        raise gr.Error(f"global exception: {error}")
-
+        raise gr.Error(f"Error during processing: {str(error)}")
 
 if __name__ == '__main__':
     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.
-    ### 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.  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! ###
     """
 
     block = gr.Blocks().queue(max_size=10)
@@ -125,7 +111,7 @@ if __name__ == '__main__':
                     ["2xRRDB", "4xRRDB", "4xGRL", "4xDAT"],
                     type="value",
                     value="4xGRL",
-                    label="model"
+                    label="Model"
                 )
                 run_btn = gr.Button(value="Submit")
 
@@ -133,20 +119,17 @@ 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])
 
-    block.launch()
+    block.launch()