import gradio as gr
import numpy as np
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.hub.snapshot_download import snapshot_download

model_dir = snapshot_download('damo/cv_unet_person-image-cartoon_compound-models', cache_dir='.')

img_cartoon = pipeline(
    Tasks.image_portrait_stylization,
    model='damo/cv_unet_person-image-cartoon_compound-models')


def infer(image):
    result = img_cartoon(image.name)
    out = result[OutputKeys.OUTPUT_IMG]
    out = np.clip(out, 0, 255).astype(np.uint8)
    return out[:, :, ::-1]


with gr.Blocks() as demo:
    with gr.Row():
        image = gr.Image(label='Input', type='file')
        run_button = gr.Button('Run')
        result = gr.Image(label='Output')
    run_button.click(fn=infer, inputs=image, outputs=result)
demo.launch()