Update app.py

app.py

@@ -1,26 +1,37 @@
 import gradio as gr
-from PIL import Image
-import torch
-from torchvision import transforms
-from transformers import pipeline
+import sys
+from pathlib import Path
+import requests
+from fastai.vision import *
+from deoldify.visualize import *
 
-# Carregar o pipeline Hugging Face para colorização
-colorization_pipeline = pipeline("image-colorization", model="huggingface/ColorfulImageColorization")
+# 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)
 
-# Função para colorizar a imagem
-def colorize_image(image):
-    image = image.convert("RGB")
-    colorized = colorization_pipeline(image)
-    return colorized
+# Configurar o modelo DeOldify
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+colorizer = get_stable_image_colorizer(root_folder=".", artistic=True)
 
-# Configuração da Interface Gradio
+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="pil", label="Upload de Imagem em Preto e Branco"),
-    outputs=gr.Image(type="pil", label="Imagem Colorida"),
+    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="Faça upload de uma imagem em preto e branco, e esta IA a colorizará automaticamente!"
+    description="Carregue uma imagem em preto e branco, e o modelo colorizará automaticamente!",
 )
 
-# Executar a aplicação
 interface.launch()