import gradio as gr
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer
from PIL import Image

device.set(device=DeviceId.CPU)  # Dùng CPU cho an toàn
colorizer = get_image_colorizer(artistic=True)

def colorize_image(image, render_factor=35):
    # Lưu tạm
    image.save("temp_input.jpg")
    
    # Tô màu
    result_path = colorizer.plot_transformed_image(
        path="temp_input.jpg",
        render_factor=render_factor,
        watermarked=False
    )
    
    return Image.open(result_path)

# Tạo giao diện Gradio
demo = gr.Interface(
    fn=colorize_image,
    inputs=[
        gr.Image(type="pil", label="Ảnh cần tô màu"),
        gr.Slider(10, 45, value=35, label="Render Factor")
    ],
    outputs=gr.Image(type="pil", label="Ảnh đã tô màu"),
    title="DeOldify - Tô màu ảnh cũ",
    description="Upload ảnh đen trắng để tô màu"
)

demo.launch()