from upcunet_v3 import RealWaifuUpScaler
import gradio as gr
import time
import logging
import os
from PIL import ImageOps
import numpy as np
import math


def greet(input_img, input_model_name, input_tile_mode):
    input_img = np.array(input_img)
    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/"
    model_cache = {}

    # 修改Gradio输入组件定义
    input_model_name = gr.Dropdown(choices=os.listdir(ModelPath), value="up2x-latest-no-denoise.pth", label='选择model')
    input_tile_mode = gr.Dropdown(choices=[0, 1, 2, 3, 4], value=2, label='选择tile_mode')
    input_img = gr.Image(label='image', type='pil')

    inputs = [input_img, input_model_name, input_tile_mode]
    outputs = gr.Image(type="numpy")  # 明确输出类型
    
    iface = gr.Interface(
        fn=greet,
        inputs=inputs,
        outputs=outputs,
        # 移除了allow_screenshot参数，该参数在Gradio 3.x中已被移除
        allow_flagging='never',
        examples=[['test-img.jpg', "up2x-latest-no-denoise.pth", 2]],
        article='[https://github.com/bilibili/ailab/tree/main/Real-CUGAN](https://github.com/bilibili/ailab/tree/main/Real-CUGAN)<br>'
                '感谢b站开源的项目，图片过大会导致内存不足，所有我将图片裁剪小，想体验大图片的效果请自行前往上面的链接。<br>'
                'The large image will lead to memory limit exceeded. So I crop and resize image. '
                'If you want to experience the large image, please go to the link above.'
    )
    iface.launch()