app.py

import gradio as gr
import torch
from transformers import (
    Blip2Processor, Blip2ForConditionalGeneration,
    AutoProcessor, AutoModelForCausalLM, AutoModelForVision2Seq
)
from PIL import Image
import numpy as np

class ModelManager:
    def __init__(self):
        self.current_model = None
        self.current_processor = None
        self.model_name = None
        
    def load_blip2(self):
        """Carrega modelo BLIP-2"""
        self.model_name = "Salesforce/blip2-opt-2.7b"
        self.current_processor = Blip2Processor.from_pretrained(self.model_name)
        self.current_model = Blip2ForConditionalGeneration.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        return "BLIP-2 carregado com sucesso!"
    
    def load_llava(self):
        """Carrega modelo LLaVA"""
        self.model_name = "llava-hf/llava-1.5-7b-hf"
        self.current_processor = AutoProcessor.from_pretrained(self.model_name)
        self.current_model = AutoModelForVision2Seq.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        return "LLaVA carregado com sucesso!"
    
    def load_git(self):
        """Carrega modelo GIT"""
        self.model_name = "microsoft/git-base-coco"
        self.current_processor = AutoProcessor.from_pretrained(self.model_name)
        self.current_model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        return "GIT carregado com sucesso!"

    def analyze_image(self, image, question, model_choice):
        """Analisa imagem com o modelo selecionado"""
        try:
            # Carrega o modelo apropriado se necessário
            if model_choice == "BLIP-2" and (self.model_name != "Salesforce/blip2-opt-2.7b"):
                status = self.load_blip2()
            elif model_choice == "LLaVA" and (self.model_name != "llava-hf/llava-1.5-7b-hf"):
                status = self.load_llava()
            elif model_choice == "GIT" and (self.model_name != "microsoft/git-base-coco"):
                status = self.load_git()
            
            # Prepara a imagem
            if isinstance(image, str):
                image = Image.open(image)
            elif isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            
            # Processa a entrada
            inputs = self.current_processor(
                images=image,
                text=question,
                return_tensors="pt"
            ).to(self.current_model.device)
            
            # Gera a resposta
            outputs = self.current_model.generate(
                **inputs,
                max_new_tokens=150,
                num_beams=5,
                temperature=0.7,
                top_p=0.9
            )
            
            response = self.current_processor.decode(outputs[0], skip_special_tokens=True)
            return response
        
        except Exception as e:
            return f"Erro na análise: {str(e)}"

# Cria instância do gerenciador de modelos
model_manager = ModelManager()

# Interface Gradio
with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("""
    # 🤖 Análise Visual Multi-Modelo
    Escolha o modelo que deseja usar para analisar sua imagem
    """)
    
    with gr.Row():
        with gr.Column():
            # Inputs
            model_choice = gr.Radio(
                choices=["BLIP-2", "LLaVA", "GIT"],
                label="Escolha o Modelo",
                value="BLIP-2"
            )
            
            with gr.Box():
                gr.Markdown("""
                ### 📝 Características dos Modelos:
                
                **BLIP-2:**
                - Melhor para análise técnica
                - Mais preciso em detalhes
                - Respostas estruturadas
                
                **LLaVA:**
                - Mais conversacional
                - Respostas naturais
                - Bom para perguntas abertas
                
                **GIT:**
                - Mais rápido e leve
                - Bom para descrições simples
                - Menor uso de memória
                """)
            
            image_input = gr.Image(
                type="pil",
                label="Imagem para Análise"
            )
            
            question_input = gr.Textbox(
                label="Sua Pergunta",
                placeholder="Faça uma pergunta sobre a imagem..."
            )
            
            analyze_btn = gr.Button("🔍 Analisar", variant="primary")
        
        with gr.Column():
            # Output
            with gr.Box():
                gr.Markdown("### 📊 Resultado da Análise")
                output_text = gr.Markdown()
            
            with gr.Accordion("💡 Sugestões de Perguntas", open=False):
                gr.Markdown("""
                1. Quais alimentos você identifica na imagem?
                2. Como os alimentos estão preparados?
                3. Descreva a apresentação do prato.
                4. Que tipo de culinária parece ser?
                5. Quais são os ingredientes principais?
                6. Como está a disposição dos elementos no prato?
                7. Há algum detalhe de decoração ou finalização?
                8. Qual parece ser o método de cocção utilizado?
                """)
    
    # Eventos
    analyze_btn.click(
        fn=model_manager.analyze_image,
        inputs=[image_input, question_input, model_choice],
        outputs=output_text
    )

if __name__ == "__main__":
    print(f"Dispositivo: {'CUDA' if torch.cuda.is_available() else 'CPU'}")
    iface.launch()