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

# Carregar o modelo
model = AutoModel.from_pretrained("openbmb/MiniCPM-o-2_6", trust_remote_code=True)

# Função para processar a imagem e gerar uma resposta
def process_image(input_image):
    # Pré-processamento da imagem (se necessário)
    # Exemplo: redimensionar, normalizar, converter para tensor, etc.
    input_image = input_image.resize((224, 224))  # Redimensionar para o tamanho esperado pelo modelo
    input_tensor = torch.tensor(np.array(input_image)).permute(2, 0, 1).unsqueeze(0).float()  # Converter para tensor

    # Passar a imagem pelo modelo
    with torch.no_grad():
        output = model(input_tensor)

    # Processar a saída do modelo (depende do que o modelo retorna)
    # Aqui estou apenas simulando uma resposta
    response = "Processamento concluído. Saída do modelo disponível."

    return response

# Criar a interface Gradio
interface = gr.Interface(
    fn=process_image,  # Função que processa a imagem
    inputs="image",    # Tipo de entrada: imagem
    outputs="text",    # Tipo de saída: texto
    title="MiniCPM-o-2_6 - Processamento de Imagens",  # Título da interface
    description="Envie uma imagem e veja a resposta do modelo MiniCPM-o-2_6."  # Descrição
)

# Iniciar a interface
interface.launch()