import os
import torch
from PIL import Image
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration

# 1. Carregar o modelo BLIP finetuned
MODEL_DIR = "Frame_30K_gpu_00_ptt"  # Substitua pelo nome correto da pasta do seu modelo

# Verifique se o diretório do modelo existe localmente

processor = BlipProcessor.from_pretrained(MODEL_DIR)
model = BlipForConditionalGeneration.from_pretrained(MODEL_DIR)

# Mover o modelo para o dispositivo disponível
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def generate_caption(image):
    """
    Gera uma legenda para a imagem fornecida usando o modelo BLIP finetuned.

    Args:
        image (PIL.Image.Image): Imagem carregada.

    Returns:
        str: Legenda gerada.
    """
    if image is None:
        return "Nenhuma imagem fornecida."

    # Preprocessar a imagem
    inputs = processor(images=image, return_tensors="pt").to(device)

    # Gerar legenda
    out = model.generate(
        **inputs,
        max_length=100,  # Ajuste conforme necessário
        num_beams=5, 
        no_repeat_ngram_size=2
    )
    caption = processor.decode(out[0], skip_special_tokens=True)
    
    return caption

# Configurar a interface Gradio
iface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Enviar Imagem"),
    outputs=gr.Textbox(
        label="Legenda Gerada",
        lines=4,               # Número de linhas visíveis
        max_lines=10,          # Número máximo de linhas permitidas
        interactive=False      # Impede a edição pelo usuário
    ),
    title="BLIP Português - Geração de Legendas para Imagens",
    description="Envie uma imagem e o modelo BLIP finetuned em português irá gerar uma legenda descritiva para ela.",
    examples=[
        ["examples/image1.jpg"],
        ["examples/image2.jpg"],
        ["examples/image3.jpg"],
    ],
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()