feat: update queue

app.py

@@ -1,5 +1,6 @@
 import os, sys
 import cv2
+import time
 import gradio as gr
 import torch
 import numpy as np
@@ -47,15 +48,16 @@ def inference(img_path, model_name):
             generator = load_rrdb(weight_path, scale=2) # Directly use default way now
             
         else:
-            raise gr.Error(error)
+            raise gr.Error("We don't support such Model")
         
         generator = generator.to(dtype=weight_dtype)
 
 
         # 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, crop_for_4x=True)
-        save_image(super_resolved_img, "SR_result.png")
-        outputs = cv2.imread("SR_result.png")
+        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)
         
         return outputs
@@ -70,14 +72,16 @@ if __name__ == '__main__':
     
     MARKDOWN = \
     """
-    ## APISR: Anime Production Inspired Real-World Anime Super-Resolution (CVPR 2024)
-
+    ## <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)
 
-    If APISR is helpful for you, please help star the GitHub Repo. Thanks!
+    APISR aims at restoring and enhancing low-quality low-resolution anime images and video sources with various degradations from real-world scenarios.
+
+    If APISR is helpful, please help star the GitHub Repo. Thanks! 
     """
 
-    block = gr.Blocks().queue()
+    block = gr.Blocks().queue(max_size=10)
     with block:
         with gr.Row():
             gr.Markdown(MARKDOWN)

test_code/inference.py

@@ -4,6 +4,7 @@
 import argparse
 import os, sys, cv2, shutil, warnings
 import torch
+import gradio as gr
 from torchvision.transforms import ToTensor
 from torchvision.utils import save_image
 warnings.simplefilter("default")
@@ -38,6 +39,12 @@ def super_resolve_img(generator, input_path, output_path=None, weight_dtype=torc
             img_lr = img_lr[:4*(h//4),:,:]
         if w % 4 != 0:
             img_lr = img_lr[:,:4*(w//4),:]
+                
+    # Check if the size is out of the boundary
+    h, w, c = img_lr.shape
+    if h*w > 720*1280:
+        raise gr.Error("The input image size is too large. The largest area we support is 720x1280=921600 pixel!")
+    
 
     # Transform to tensor
     img_lr = cv2.cvtColor(img_lr, cv2.COLOR_BGR2RGB)
@@ -54,7 +61,7 @@ def super_resolve_img(generator, input_path, output_path=None, weight_dtype=torc
         if output_path is not None:
             save_image(super_resolved_img, output_path)
 
-    # Empty the cache everytime you finish processing one image
+    # Empty the cache every time you finish processing one image
     torch.cuda.empty_cache() 
     
     return super_resolved_img
@@ -71,7 +78,7 @@ if __name__ == "__main__":
     parser.add_argument('--scale', type = int, default = 4, help="Up scaler factor")
     parser.add_argument('--weight_path', type = str, default = 'pretrained/4x_APISR_GRL_GAN_generator.pth', help="Weight path directory, usually under saved_models folder")
     parser.add_argument('--store_dir', type = str, default = 'sample_outputs', help="The folder to store the super-resolved images")
-    parser.add_argument('--float16_inference', type = bool, default = False, help="The folder to store the super-resolved images")      # Currently, this is only supported in RRDB, there is some bug with GRL model
+    parser.add_argument('--float16_inference', type = bool, default = False, help="Float16 inference, only useful in RRDB now")      # Currently, this is only supported in RRDB, there is some bug with GRL model
     args = parser.parse_args()
     
     # Sample Command