Update app.py

app.py

@@ -1,36 +1,73 @@
-from transformers import AutoModel
-import gradio as gr
-from PIL import Image
 import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+import gradio as gr
+
+# Configuração inicial
+torch.manual_seed(100)
 
-# Carregar o modelo
-model = AutoModel.from_pretrained("openbmb/MiniCPM-o-2_6", trust_remote_code=True)
+# Carregar o modelo e o tokenizer
+model = AutoModel.from_pretrained('openbmb/MiniCPM-o-2_6', trust_remote_code=True,
+                                  attn_implementation='sdpa', torch_dtype=torch.bfloat16)
+model = model.eval().cuda()
+tokenizer = AutoTokenizer.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
+# Função para interagir com o modelo
+def chat_with_model(image, question, chat_history=None):
+    if chat_history is None:
+        chat_history = []
 
-    # Passar a imagem pelo modelo
-    with torch.no_grad():
-        output = model(input_tensor)
+    # Converter a imagem para RGB (se necessário)
+    if isinstance(image, str):
+        image = Image.open(image).convert('RGB')
+    else:
+        image = image.convert('RGB')
 
-    # 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."
+    # Preparar a mensagem para o modelo
+    msgs = [{'role': 'user', 'content': [image, question]}]
 
-    return response
+    # Adicionar histórico de conversa, se houver
+    for msg in chat_history:
+        msgs.append(msg)
+
+    # Gerar resposta do modelo
+    answer = model.chat(
+        msgs=msgs,
+        tokenizer=tokenizer
+    )
+
+    # Atualizar o histórico de conversa
+    chat_history.append({"role": "user", "content": [image, question]})
+    chat_history.append({"role": "assistant", "content": [answer]})
+
+    # Retornar a resposta e o histórico atualizado
+    return answer, chat_history
+
+# Interface Gradio
+def gradio_interface(image, question, chat_history=None):
+    response, updated_history = chat_with_model(image, question, chat_history)
+    return response, updated_history
 
 # 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
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# MiniCPM-o-2_6 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")
+        question_input = gr.Textbox(label="Your Question", placeholder="What is in the image?")
+    
+    chat_history = gr.State([])  # Armazenar o histórico de conversa
+    output_text = gr.Textbox(label="Model Response", interactive=False)
+    
+    submit_button = gr.Button("Submit")
+    
+    # Ação ao clicar no botão
+    submit_button.click(
+        fn=gradio_interface,
+        inputs=[image_input, question_input, chat_history],
+        outputs=[output_text, chat_history]
+    )
 
 # Iniciar a interface
-interface.launch()
+demo.launch()