import os 
os.system('pip install "modelscope[cv]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html')
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:
    title= gr.Markdown("""
    # Gradio Demo for [DCT-Net: Domain-Calibrated Translation for Portrait Stylization](https://github.com/menyifang/DCT-Net), SIGGRAPH 2022 (TOG); Multi-style cartoonization
    """
    )
    with gr.Row():
        image = gr.Image(label='Input', type='file')
        result = gr.Image(label='Output')
    run_button = gr.Button('Run')
    run_button.click(fn=infer, inputs=image, outputs=result)
demo.launch()