import os
from pathlib import Path

import gradio as gr
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.generators import gen_inference_deep
from huggingface_hub import snapshot_download

from model_image_colorizer import ImageFilter, ModelImageColorizer

os.system("pip freeze")

device.set(device=DeviceId.CPU)

REPO_ID = "leonelhs/deoldify"
MODEL_NAME = "ColorizeArtistic_gen"

snapshot_folder = snapshot_download(repo_id=REPO_ID)
learn = gen_inference_deep(root_folder=Path(snapshot_folder), weights_name=MODEL_NAME)
image_filter = ImageFilter(learn=learn)
colorizer = ModelImageColorizer(image_filter)


def predict(image):
    return colorizer.get_colored_image(image, render_factor=35)


title = "DeOldify"
description = r"""
## Colorize image

This is an implementation of <a href='https://github.com/jantic/DeOldify' target='_blank'>DeOldify</a>.
It has no any particular purpose than start research on AI models.

"""

article = r"""
Questions, doubts, comments, please email 📧 `leonelhs@gmail.com`

This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/jantic/DeOldify' target='_blank'>Github ⭐</a>

<a href="https://www.buymeacoffee.com/leonelhs"><img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" /></a>

<center><img src='https://visitor-badge.glitch.me/badge?page_id=deoldify.visitor-badge' alt='visitor badge'></center>
"""

demo = gr.Interface(
    predict, [
        gr.Image(type="pil", label="Image gray scale"),
    ], [
        gr.Image(type="pil", label="Image color")
    ],
    title=title,
    description=description,
    article=article)

demo.queue().launch()