import gradio as gr
import sys
from pathlib import Path
import requests
from fastai.vision import *
from deoldify.visualize import *

# Baixar pesos do modelo
model_path = Path("./models/ColorizeArtistic_gen.pth")
if not model_path.exists():
    model_path.parent.mkdir(parents=True, exist_ok=True)
    url = "https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth"
    response = requests.get(url, stream=True)
    with open(model_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=1024):
            f.write(chunk)

# Configurar o modelo DeOldify
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
colorizer = get_stable_image_colorizer(root_folder=".", artistic=True)

def colorize_image(input_image):
    output = colorizer.plot_transformed_image(
        path=input_image, render_factor=35, display_render_factor=True, figsize=(20, 20)
    )
    return output

# Interface do Gradio
interface = gr.Interface(
    fn=colorize_image,
    inputs=gr.Image(type="filepath", label="Imagem em Preto e Branco"),
    outputs=gr.Image(type="auto", label="Imagem Colorida"),
    title="Colorização de Imagens com IA",
    description="Carregue uma imagem em preto e branco, e o modelo colorizará automaticamente!",
)

interface.launch()