Update app.py

app.py

@@ -1,180 +1,249 @@
 import gradio as gr
 import torch
-from transformers import (
-    Blip2Processor, Blip2ForConditionalGeneration,
-    AutoProcessor, AutoModelForCausalLM, AutoModelForVision2Seq
-)
+from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
 import numpy as np
+import os
+from huggingface_hub import snapshot_download
+import logging
 
-class ModelManager:
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class NutritionalAnalyzer:
     def __init__(self):
-        self.current_model = None
-        self.current_processor = None
-        self.model_name = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.models = {}
+        self.processors = {}
         
-    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 initialize_model(self, model_name):
+        """Initialize a specific model with error handling and caching"""
+        try:
+            if model_name not in self.models:
+                logger.info(f"Initializing {model_name}...")
+                
+                # Model-specific configurations
+                model_configs = {
+                    "llava": {
+                        "repo": "llava-hf/llava-1.5-7b-hf",
+                        "local_cache": "models/llava"
+                    },
+                    "git": {
+                        "repo": "microsoft/git-base-coco",
+                        "local_cache": "models/git"
+                    }
+                }
+                
+                config = model_configs.get(model_name)
+                if not config:
+                    raise ValueError(f"Unsupported model: {model_name}")
+                
+                # Ensure cache directory exists
+                os.makedirs(config["local_cache"], exist_ok=True)
+                
+                # Download model if not cached
+                if not os.path.exists(os.path.join(config["local_cache"], "model.safetensors")):
+                    snapshot_download(
+                        repo_id=config["repo"],
+                        local_dir=config["local_cache"],
+                        ignore_patterns=["*.md", "*.txt"]
+                    )
+                
+                # Load processor and model
+                self.processors[model_name] = AutoProcessor.from_pretrained(
+                    config["local_cache"],
+                    local_files_only=True
+                )
+                
+                self.models[model_name] = AutoModelForVision2Seq.from_pretrained(
+                    config["local_cache"],
+                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                    device_map="auto",
+                    local_files_only=True
+                )
+                
+                logger.info(f"{model_name} initialized successfully")
+                return True
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error initializing {model_name}: {str(e)}")
+            return False
+
+    def prepare_image(self, image):
+        """Prepare image for model input"""
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        # Ensure image is in RGB mode
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+            
+        return image
+
+    def generate_nutritional_prompt(self, user_question):
+        """Generate a comprehensive nutritional analysis prompt"""
+        return f"""Como nutricionista especializado, analise esta refeição detalhadamente:
+
+1. Composição do Prato:
+   - Ingredientes principais
+   - Proporções aproximadas
+   - Método de preparo aparente
+
+2. Análise Nutricional:
+   - Estimativa calórica
+   - Macronutrientes (proteínas, carboidratos, gorduras)
+   - Principais micronutrientes
+
+3. Recomendações:
+   - Sugestões para versão mais saudável
+   - Porção recomendada
+   - Adequação para dietas específicas
+
+Pergunta específica do usuário: {user_question}
+
+Por favor, forneça uma análise detalhada em português."""
 
     def analyze_image(self, image, question, model_choice):
-        """Analisa imagem com foco nutricional"""
+        """Analyze image with nutritional focus"""
         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()
+            # Convert model choice to internal name
+            model_name = model_choice.lower().replace("-", "")
             
-            # Adiciona contexto nutricional à pergunta
-            nutritional_prompt = (
-                "Como nutricionista, analise este prato considerando: "
-                "1. Lista de ingredientes principais\n"
-                "2. Estimativa calórica total\n"
-                "3. Sugestões para uma versão mais saudável\n"
-                "4. Análise de grupos alimentares\n"
-                f"Pergunta do usuário: {question}"
-                "\nPor favor, responda em português com detalhes nutricionais."
-            )
-
-            # Prepara a imagem
-            if isinstance(image, str):
-                image = Image.open(image)
-            elif isinstance(image, np.ndarray):
-                image = Image.fromarray(image)
+            # Initialize model if needed
+            if not self.initialize_model(model_name):
+                return "Erro: Não foi possível inicializar o modelo. Por favor, tente novamente."
             
-            # Processa a entrada
-            inputs = self.current_processor(
-                images=image,
+            # Prepare image and prompt
+            processed_image = self.prepare_image(image)
+            nutritional_prompt = self.generate_nutritional_prompt(question)
+            
+            # Process input
+            inputs = self.processors[model_name](
+                images=processed_image,
                 text=nutritional_prompt,
                 return_tensors="pt"
-            ).to(self.current_model.device)
+            ).to(self.device)
             
-            # Gera a resposta
-            outputs = self.current_model.generate(
-                **inputs,
-                max_new_tokens=200,  # Aumentado para respostas mais completas
-                num_beams=5,
-                temperature=0.7,
-                top_p=0.9
-            )
+            # Generate response with enhanced parameters
+            with torch.no_grad():
+                outputs = self.models[model_name].generate(
+                    **inputs,
+                    max_new_tokens=300,
+                    num_beams=5,
+                    temperature=0.7,
+                    top_p=0.9,
+                    repetition_penalty=1.2,
+                    length_penalty=1.0
+                )
+            
+            # Decode and format response
+            response = self.processors[model_name].decode(outputs[0], skip_special_tokens=True)
+            formatted_response = self.format_response(response)
+            
+            return formatted_response
             
-            # Decodifica e formata a resposta
-            response = self.current_processor.decode(outputs[0], skip_special_tokens=True)
-            formatted_response = response.replace(". ", ".\n").replace("; ", ";\n")
-            return f"**Análise Nutricional:**\n{formatted_response}"
-        
         except Exception as e:
-            return f"Erro na análise: {str(e)}"
+            logger.error(f"Analysis error: {str(e)}")
+            return f"Erro na análise: {str(e)}\nPor favor, tente novamente ou escolha outro modelo."
 
-# Cria instância do gerenciador de modelos
-model_manager = ModelManager()
+    def format_response(self, response):
+        """Format the response for better readability"""
+        sections = [
+            "Composição do Prato",
+            "Análise Nutricional",
+            "Recomendações"
+        ]
+        
+        formatted = "# 📊 Análise Nutricional\n\n"
+        
+        # Split response into paragraphs
+        paragraphs = response.split("\n")
+        
+        current_section = ""
+        for paragraph in paragraphs:
+            # Check if paragraph starts a new section
+            for section in sections:
+                if section.lower() in paragraph.lower():
+                    current_section = f"\n## {section}\n"
+                    formatted += current_section
+                    break
+            
+            # Add paragraph to current section
+            if paragraph.strip() and current_section:
+                formatted += f"- {paragraph.strip()}\n"
+            elif paragraph.strip():
+                formatted += f"{paragraph.strip()}\n"
+        
+        return formatted
 
-# Interface Gradio
-with gr.Blocks(theme=gr.themes.Soft()) as iface:
-    gr.Markdown("""
-    # 🥗 Analisador Nutricional Inteligente
-    Escolha o modelo que deseja usar para analisar seu prato e obter recomendações nutricionais.
-    """)
+# Create interface
+def create_interface():
+    analyzer = NutritionalAnalyzer()
     
-    with gr.Row():
-        with gr.Column():
-            # Inputs
-            model_choice = gr.Radio(
-                choices=["BLIP-2", "LLaVA", "GIT"],
-                label="Escolha o Modelo",
-                value="BLIP-2"
-            )
-            
-            # Substitui gr.Box() por gr.Group() para compatibilidade
-            with gr.Group():
-                gr.Markdown("""
-                ### 📝 Características dos Modelos:
+    with gr.Blocks(theme=gr.themes.Soft()) as iface:
+        gr.Markdown("""
+        # 🥗 Análise Nutricional Inteligente
+        Faça upload da foto do seu prato para receber uma análise nutricional detalhada.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                image_input = gr.Image(
+                    type="pil",
+                    label="📸 Foto do Prato",
+                    height=400
+                )
                 
-                **BLIP-2:**
-                - Análise detalhada de ingredientes
-                - Estimativas calóricas mais precisas
-                - Recomendações técnicas
+                question_input = gr.Textbox(
+                    label="💭 Sua Pergunta",
+                    placeholder="Ex: Quais são os nutrientes principais deste prato?",
+                    lines=2
+                )
                 
-                **LLaVA:**
-                - Explicações mais conversacionais
-                - Sugestões práticas para o dia a dia
-                - Foco em hábitos alimentares
+                model_choice = gr.Radio(
+                    choices=["LLaVA", "GIT"],
+                    value="LLaVA",
+                    label="🤖 Escolha o Modelo de Análise"
+                )
                 
-                **GIT:**
-                - Respostas rápidas e diretas
-                - Ideal para análises simples
-                - Menor consumo de recursos
-                """)
-            
-            image_input = gr.Image(
-                type="pil",
-                label="Foto do Prato"
-            )
-            
-            question_input = gr.Textbox(
-                label="Sua Pergunta",
-                placeholder="Ex: Quantas calorias tem este prato? Como posso torná-lo mais saudável?"
-            )
+                analyze_btn = gr.Button(
+                    "🔍 Analisar Prato",
+                    variant="primary",
+                    scale=1
+                )
             
-            analyze_btn = gr.Button("🔍 Analisar", variant="primary")
+            with gr.Column(scale=3):
+                output = gr.Markdown(label="Resultado da Análise")
         
-        with gr.Column():
-            # Output
-            with gr.Group():  # Substitui gr.Box() por gr.Group()
-                gr.Markdown("### 📊 Resultado da Análise")
-                output_text = gr.Markdown()
+        # Add examples and tips
+        with gr.Accordion("💡 Dicas de Uso", open=False):
+            gr.Markdown("""
+            ### Sugestões de Perguntas:
+            - Qual o valor nutricional aproximado deste prato?
+            - Como tornar esta refeição mais equilibrada?
+            - Este prato é adequado para dieta low-carb?
+            - Quais nutrientes importantes estão presentes?
             
-            with gr.Accordion("💡 Sugestões de Perguntas", open=False):
-                gr.Markdown("""
-                1. Quantas calorias tem este prato?
-                2. Quais são os ingredientes principais?
-                3. Como posso tornar este prato mais saudável?
-                4. Este prato é adequado para uma dieta low-carb?
-                5. Quais nutrientes estão presentes neste prato?
-                6. Este prato é rico em proteínas?
-                7. Como posso substituir ingredientes para reduzir calorias?
-                8. Este prato é indicado para quem tem restrição a glúten/lactose?
-                """)
+            ### Dicas para Melhores Resultados:
+            1. Tire a foto com boa iluminação
+            2. Capture todos os elementos do prato
+            3. Evite ângulos muito inclinados
+            4. Seja específico em suas perguntas
+            """)
+        
+        analyze_btn.click(
+            fn=analyzer.analyze_image,
+            inputs=[image_input, question_input, model_choice],
+            outputs=output
+        )
     
-    # Eventos
-    analyze_btn.click(
-        fn=model_manager.analyze_image,
-        inputs=[image_input, question_input, model_choice],
-        outputs=output_text
-    )
+    return iface
 
 if __name__ == "__main__":
-    print(f"Dispositivo: {'CUDA' if torch.cuda.is_available() else 'CPU'}")
+    iface = create_interface()
     iface.launch()