add more model choice

app.py

@@ -1,22 +1,49 @@
-import cv2
 from upcunet_v3 import RealWaifuUpScaler
-import traceback
 import gradio as gr
+import time
+import logging
+import os
 
 
-def greet(input_img):
-    print(input_img.shape)
-    result = upscaler(input_img, tile_mode=2)
-    print(result.shape)
+def greet(input_img, input_model_name, input_tile_mode):
+    if input_model_name not in model_cache:
+        t1 = time.time()
+        upscaler = RealWaifuUpScaler(input_model_name[2], ModelPath + input_model_name, half=False, device="cpu")
+        t2 = time.time()
+        logger.info(f'load model time, {t2 - t1}')
+        model_cache[input_model_name] = upscaler
+    else:
+        upscaler = model_cache[input_model_name]
+        logger.info(f'load model from cache')
+
+    start = time.time()
+    result = upscaler(input_img, tile_mode=input_tile_mode)
+    end = time.time()
+    logger.info(f'input_model_name, {input_model_name}')
+    logger.info(f'input_tile_mode, {input_tile_mode}')
+    logger.info(f'input shape, {input_img.shape}')
+    logger.info(f'output shape, {result.shape}')
+    logger.info(f'speed time, {end - start}')
     return result
 
 
 if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(process)d] [%(levelname)s] %(message)s")
+    logger = logging.getLogger()
+
     ModelPath = "weights_v3/"
-    ModelName = "up2x-latest-no-denoise.pth"
-    upscaler = RealWaifuUpScaler(2, ModelPath + ModelName, half=False, device="cpu")
+    model_cache = {}
+
+    input_model_name = gr.inputs.Dropdown(os.listdir(ModelPath), default="up2x-latest-no-denoise.pth", label='选择model')
+    input_tile_mode = gr.inputs.Dropdown([0, 1, 2, 3, 4], default=2, label='选择tile_mode')
+    input_img = gr.inputs.Image(label='image')
 
-    inputs = gr.inputs.Image()
+    inputs = [input_img, input_model_name, input_tile_mode]
     outputs = "image"
-    iface = gr.Interface(fn=greet, inputs=inputs, outputs=outputs)
+    iface = gr.Interface(fn=greet,
+                         inputs=inputs,
+                         outputs=outputs,
+                         allow_screenshot=False,
+                         allow_flagging='never',
+                         description='https://github.com/bilibili/ailab/tree/main/Real-CUGAN\n感谢b站开源的项目')
     iface.launch()