import warnings
from pathlib import Path

import gradio as gr
from huggingface_hub import snapshot_download
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer


warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")


try:
    device.set(device=DeviceId.GPU0)
except Exception:
    device.set(device=DeviceId.CPU)


REPO_ID = "leonelhs/deoldify" 
snapshot_folder = snapshot_download(repo_id=REPO_ID)


colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)


def predict(image, render_factor=35):
    return colorizer.get_transformed_image(image, render_factor=render_factor, watermarked=False)


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="filepath", label="Upload ảnh trắng đen"),
        gr.Slider(7, 40, value=35, step=1, label="Render Factor (Chất lượng)")
    ],
    outputs=gr.Image(type="pil", label="Ảnh sau khi tô màu"),
    title="Phục dựng ảnh lịch sử",
    description="Upload ảnh lịch sử.",
)

if __name__ == "__main__":
    demo.queue().launch()