import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
import gradio as gr

# Verificar se a GPU está disponível
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else torch.float16

# Carregar o modelo e o tokenizer
model = AutoModel.from_pretrained('openbmb/MiniCPM-V', trust_remote_code=True, torch_dtype=dtype)
model = model.to(device=device, dtype=dtype)  # Ajuste para o dispositivo e tipo de dados adequados
tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V', trust_remote_code=True)
model.eval()

# Função para processar a imagem e a pergunta
def chat_with_model(image, question):
    try:
        # Converter a imagem para RGB (se necessário)
        if isinstance(image, str):
            image = Image.open(image).convert('RGB')
        else:
            image = image.convert('RGB')

        # Preparar a mensagem para o modelo
        msgs = [{'role': 'user', 'content': question}]

        # Gerar resposta do modelo
        res, context, _ = model.chat(
            image=image,
            msgs=msgs,
            context=None,
            tokenizer=tokenizer,
            sampling=True,
            temperature=0.7
        )

        return res
    except Exception as e:
        return f"Erro ao processar a imagem ou pergunta: {str(e)}"

# Interface Gradio
def gradio_interface(image, question):
    response = chat_with_model(image, question)
    return response

# Criar a interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# MiniCPM-V Chat with Images")
    gr.Markdown("Envie uma imagem e faça perguntas sobre ela.")

    with gr.Row():
        image_input = gr.Image(label="Upload Image", type="pil")  # Campo para upload de imagem
        question_input = gr.Textbox(label="Your Question", placeholder="What is in the image?")  # Campo para a pergunta
    
    output_text = gr.Textbox(label="Model Response", interactive=False)  # Campo para exibir a resposta
    
    submit_button = gr.Button("Submit")  # Botão para enviar
    
    # Ação ao clicar no botão
    submit_button.click(
        fn=gradio_interface,
        inputs=[image_input, question_input],
        outputs=output_text
    )

# Iniciar a interface
demo.launch()