Arcadia/clarify

1from upcunet_v3 import RealWaifuUpScaler2import gradio as gr3import time4import logging5import os6from PIL import ImageOps7import numpy as np8import math9 10 11 12def greet(input_img, input_model_name, input_tile_mode):13    # if input_img.size[0] * input_img.size[1] > 256 * 256:14    #     y = int(math.sqrt(256*256/input_img.size[0]*input_img.size[1]))15    #     x = int(input_img.size[0]/input_img.size[1]*y)16    #     input_img = ImageOps.fit(input_img, (x, y))17    input_img = np.array(input_img)18    if input_model_name not in model_cache:19        t1 = time.time()20        upscaler = RealWaifuUpScaler(input_model_name[2], ModelPath + input_model_name, half=False, device="cpu")21        t2 = time.time()22        logger.info(f'load model time, {t2 - t1}')23        model_cache[input_model_name] = upscaler24    else:25        upscaler = model_cache[input_model_name]26        logger.info(f'load model from cache')27 28    start = time.time()29    result = upscaler(input_img, tile_mode=input_tile_mode)30    end = time.time()31    logger.info(f'input_model_name, {input_model_name}')32    logger.info(f'input_tile_mode, {input_tile_mode}')33    logger.info(f'input shape, {input_img.shape}')34    logger.info(f'output shape, {result.shape}')35    logger.info(f'speed time, {end - start}')36    return result37 38 39if __name__ == '__main__':40    logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(process)d] [%(levelname)s] %(message)s")41    logger = logging.getLogger()42 43    ModelPath = "weights_v3/"44    model_cache = {}45 46    input_model_name = gr.inputs.Dropdown(os.listdir(ModelPath), default="up2x-latest-denoise2x.pth", label='选择model')47    input_tile_mode = gr.inputs.Dropdown([0, 1, 2, 3, 4], default=2, label='选择tile_mode')48    input_img = gr.inputs.Image(label='image', type='pil')49 50    inputs = [input_img, input_model_name, input_tile_mode]51    outputs = "image"52    iface = gr.Interface(fn=greet,53                         inputs=inputs,54                         outputs=outputs,55                         allow_screenshot=False,56                         allow_flagging='never',57                         examples=[['test-img.jpg', "up2x-latest-denoise2x.pth", 2]],58                         article=' ')59    iface.launch()60