import os

os.system('wget "https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth" -O "models/ColorizeArtistic_gen.pth"')

import gradio as gr
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
import tempfile

# if not os.path.exists("models/ColorizeArtistic_gen.pth"):
    

device.set(device=DeviceId.GPU0)

colorizer = get_image_colorizer(artistic=True)


def colorize(image):
    tmp_folder = tempfile.TemporaryDirectory()
    return colorizer.plot_transformed_image_from_url(
        url=image,
        path=f"{tmp_folder.name}/input.png",
        render_factor=35,
        compare=True,
        results_dir=f"{tmp_folder.name}/output.png")


gr.Interface(colorize, ["text"],
             ["image"]).queue(default_enabled=True).launch(show_api=True)