Update app.py

app.py

@@ -5,45 +5,60 @@ from typing import Optional
 from PIL import Image
 import base64
 import io
+import requests
 from dotenv import load_dotenv
 from functools import lru_cache
 import torch
 
-# ------------ CONFIGURACIÓN ------------
-load_dotenv()
-
+# ------------ CONFIGURACIÓN MEJORADA ------------
 class Config:
     def __init__(self):
+        # PRIMERO busca en secrets de Hugging Face, LUEGO en .env
         self.SAMBANOVA_API_KEY = os.getenv("SAMBANOVA_API_KEY")
         self.BRIA_API_TOKEN = os.getenv("BRIA_API_TOKEN")
         self.HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+        
+        # Si no está en environment, intenta cargar desde huggingface_hub
+        if not self.SAMBANOVA_API_KEY:
+            try:
+                from huggingface_hub import HfFolder
+                self.SAMBANOVA_API_KEY = HfFolder.get_token()
+            except ImportError:
+                pass
+        
         self.validate_keys()
 
     def validate_keys(self):
         if not self.SAMBANOVA_API_KEY:
-            raise ValueError("❌ Falta SAMBANOVA_API_KEY en .env")
+            print("❌ SAMBANOVA_API_KEY no encontrado")
+            print("💡 Configura el secret en Hugging Face Spaces:")
+            print("   - Ve a Settings → Repository secrets")
+            print("   - Agrega: SAMBANOVA_API_KEY = tu_api_key_real")
+        else:
+            print(f"✅ SAMBANOVA_API_KEY configurado (primeros 10 chars): {self.SAMBANOVA_API_KEY[:10]}...")
+            
         if not self.BRIA_API_TOKEN:
-            print("⚠️ BRIA_API_TOKEN no encontrado - Funciones de imagen desactivadas")
+            print("⚠️ BRIA_API_TOKEN no configurado - Imagen deshabilitada")
+        else:
+            print("✅ BRIA_API_TOKEN configurado")
 
 config = Config()
 
 # ------------ IMPORTACION CONDICIONAL ------------
 try:
     from sambanova import SambaNova
-except ImportError:
+    print("✅ SambaNova importado correctamente")
+except ImportError as e:
+    print(f"❌ Error importando SambaNova: {e}")
+    print("💡 Instala con: pip install sambanova")
     SambaNova = None
 
-try:
-    from mcp_use import MCPClient
-except ImportError:
-    MCPClient = None
-
 # ------------ MODELOS INTEGRADOS ------------
 MODELS = {
     "chat": {
         "llama3": "NousResearch/Meta-Llama-3-8B-Instruct",
         "qwen": "Qwen/Qwen1.5-32B-Chat",
-        "deepseek": "deepseek-ai/deepseek-v3",
+        "deepseek": "deepseek-ai/deepseek-llm-67b-chat",
         "sambanova": {
             "models": ["Llama-3.3-Swallow-70B-Instruct-v0.4", "Qwen3-32B", "DeepSeek-V3-0324"],
             "api_key": config.SAMBANOVA_API_KEY
@@ -52,9 +67,6 @@ MODELS = {
     "code": {
         "starcoder2": "bigcode/starcoder2-7b",
         "deepseek-coder": "deepseek-ai/deepseek-coder-33b-instruct"
-    },
-    "image": {
-        "bria": config.BRIA_API_TOKEN
     }
 }
 
@@ -63,144 +75,228 @@ class ModelLoader:
     @lru_cache(maxsize=3)
     def load_hf_model(self, model_name: str):
         from transformers import AutoModelForCausalLM, AutoTokenizer
-        print(f"⚙️ Cargando modelo cuantizado int8 para CPU: {model_name}...")
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=config.HUGGINGFACE_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            use_auth_token=config.HUGGINGFACE_TOKEN,
-            load_in_8bit=True,
-            device_map={"": "cpu"},
-            low_cpu_mem_usage=True
-        )
-        model.eval()
-        return model, tokenizer
+        
+        # Usar token en lugar de use_auth_token (deprecado)
+        token = config.HUGGINGFACE_TOKEN
+        
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
+            
+            # Configuración mejorada de dispositivo
+            if torch.cuda.is_available():
+                device_map = "auto"
+                print(f"🚀 Usando GPU para {model_name}")
+            else:
+                device_map = {"": "cpu"}
+                print(f"⚡ Usando CPU para {model_name}")
+            
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                token=token,
+                load_in_8bit=True,
+                device_map=device_map,
+                low_cpu_mem_usage=True,
+                torch_dtype=torch.float16
+            )
+            model.eval()
+            return model, tokenizer
+            
+        except Exception as e:
+            print(f"❌ Error cargando modelo {model_name}: {e}")
+            return None, None
 
 model_loader = ModelLoader()
 
-# ------------ HERRAMIENTAS ------------
+# ------------ HERRAMIENTAS MEJORADAS ------------
 class AI_Tools:
     def __init__(self):
-        if SambaNova is not None:
-            self.sn_client = SambaNova(
-                api_key=config.SAMBANOVA_API_KEY,
-                base_url="https://api.sambanova.ai/v1"
-            )
+        self.sn_client = None
+        if SambaNova is not None and config.SAMBANOVA_API_KEY:
+            try:
+                self.sn_client = SambaNova(
+                    api_key=config.SAMBANOVA_API_KEY,
+                    base_url="https://api.sambanova.ai/v1"
+                )
+                print("✅ Cliente SambaNova inicializado correctamente")
+            except Exception as e:
+                print(f"❌ Error inicializando SambaNova: {e}")
         else:
-            self.sn_client = None
+            print("❌ SambaNova no disponible - verifica API key")
 
     async def generate_text(self, model_type: str, model_name: str, prompt: str) -> str:
-        if model_type == "sambanova":
-            if self.sn_client is None:
-                return "❌ SambaNova client no disponible"
-            return await self._generate_sambanova(model_name, prompt)
-        else:
-            return await self._generate_hf(model_name, prompt)
+        try:
+            if model_type == "sambanova":
+                if self.sn_client is None:
+                    return "❌ Cliente SambaNova no disponible. Verifica tu API key."
+                return await self._generate_sambanova(model_name, prompt)
+            else:
+                return await self._generate_hf(model_name, prompt)
+        except Exception as e:
+            return f"❌ Error en generación: {str(e)}"
 
     async def _generate_sambanova(self, model_name: str, prompt: str) -> str:
-        response = self.sn_client.chat.completions.create(
-            model=model_name,
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.7,
-            top_p=0.9
-        )
-        return response.choices[0].message.content
+        try:
+            # CORREGIDO: Manejo adecuado de llamadas async
+            response = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.sn_client.chat.completions.create(
+                    model=model_name,
+                    messages=[{"role": "user", "content": prompt}],
+                    temperature=0.7,
+                    top_p=0.9,
+                    max_tokens=500
+                )
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return f"❌ Error SambaNova API: {str(e)}"
 
     async def _generate_hf(self, model_name: str, prompt: str) -> str:
-        model, tokenizer = model_loader.load_hf_model(model_name)
-        inputs = tokenizer(prompt, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model.generate(**inputs, max_new_tokens=256)
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        try:
+            model, tokenizer = model_loader.load_hf_model(model_name)
+            if model is None or tokenizer is None:
+                return "❌ Error cargando modelo local"
+                
+            inputs = tokenizer(prompt, return_tensors="pt")
+            
+            # Mover inputs al dispositivo del modelo
+            if hasattr(model, 'device'):
+                inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs, 
+                    max_new_tokens=256,
+                    do_sample=True,
+                    temperature=0.7,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+            
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Remover el prompt del response
+            return response[len(prompt):].strip()
+        except Exception as e:
+            return f"❌ Error en modelo local: {str(e)}"
 
     async def generate_image(self, prompt: str) -> Optional[Image.Image]:
-        if not config.BRIA_API_TOKEN or MCPClient is None:
+        if not config.BRIA_API_TOKEN:
             return None
-
-        mcp_config = {
-            "mcpServers": {
-                "bria": {
-                    "url": "https://mcp.prod.bria-api.com/mcp",
-                    "headers": {"api_token": config.BRIA_API_TOKEN}
-                }
-            }
+        
+        url = "https://api.bria.ai/v1/generate"
+        headers = {
+            "Authorization": f"Bearer {config.BRIA_API_TOKEN}",
+            "Content-Type": "application/json"
+        }
+        json_data = {
+            "prompt": prompt,
+            "options": {"resolution": "512x512"}
         }
-        async with MCPClient.from_dict(mcp_config) as client:
-            result = await client.call_tool(
-                name="generate_image",
-                arguments={
-                    "prompt": prompt,
-                    "aspect_ratio": "9:16",
-                    "quality": "ultra_detailed"
-                }
+
+        try:
+            response = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: requests.post(url, headers=headers, json=json_data, timeout=60)
             )
-            if "base64_image" in result:
-                return Image.open(io.BytesIO(base64.b64decode(result["base64_image"])))
-        return None
+            response.raise_for_status()
+            data = response.json()
+            
+            if "image_base64" in data:
+                img_bytes = base64.b64decode(data["image_base64"])
+                return Image.open(io.BytesIO(img_bytes))
+            else:
+                return None
+        except Exception as e:
+            print(f"Error generando imagen BRIA: {e}")
+            return None
 
-# ------------ INTERFAZ ------------
+# ------------ INTERFAZ MEJORADA ------------
 tools = AI_Tools()
 
-async def process_input(
-    message: str,
-    image: Optional[Image.Image],
-    model_choice: str,
-    history: list
-) -> list:
+async def process_input(message: str, image: Optional[Image.Image], history: list) -> list:
     try:
-        if image or any(kw in message.lower() for kw in ["imagen", "genera imagen"]):
-            if model_choice == "bria" and config.BRIA_API_TOKEN and MCPClient is not None:
-                img = await tools.generate_image(message)
-                if img:
-                    return history + [(message, ("Imagen generada:", img))]
-                return history + [(message, "❌ Error generando imagen")]
-            return history + [(message, "❌ BRIA no configurado o modelo no soportado")]
-
-        elif any(kw in message.lower() for kw in ["código", "code", "programa"]):
-            model_name = MODELS["code"].get(model_choice, "starcoder2")
+        msg_lower = message.lower()
+        
+        # Detección de generación de imágenes
+        if image or any(k in msg_lower for k in ["imagen", "genera imagen", "foto", "dibujo", "picture"]):
+            img = await tools.generate_image(message)
+            if img:
+                return history + [(message, ("Imagen generada:", img))]
+            return history + [(message, "❌ Error generando imagen o BRIA no configurado")]
+
+        # Detección de generación de código
+        elif any(k in msg_lower for k in ["código", "code", "programa", "script", "función", "function"]):
+            model_name = MODELS["code"]["starcoder2"]
             response = await tools.generate_text("hf", model_name, message)
-            return history + [(message, f"``````")]
+            # CORREGIDO: Formato adecuado para código
+            return history + [(message, f"```python\n{response}\n```")]
 
+        # Chat normal - prioridad a SambaNova
         else:
-            if model_choice in MODELS["chat"]["sambanova"]["models"]:
-                response = await tools.generate_text("sambanova", model_choice, message)
-            else:
-                model_name = MODELS["chat"].get(model_choice, "llama3")
+            response = ""
+            sambanova_used = False
+            
+            # Intentar SambaNova primero si está disponible
+            if tools.sn_client is not None:
+                for smodel in MODELS["chat"]["sambanova"]["models"]:
+                    if smodel.lower() in msg_lower:
+                        response = await tools.generate_text("sambanova", smodel, message)
+                        sambanova_used = True
+                        break
+            
+            # Fallback a modelo local
+            if not sambanova_used:
+                model_name = MODELS["chat"]["llama3"]
                 response = await tools.generate_text("hf", model_name, message)
+            
             return history + [(message, response)]
 
     except Exception as e:
         return history + [(message, f"❌ Error: {str(e)}")]
 
-with gr.Blocks(title="MultiModel AI Assistant (CPU + INT8)", theme=gr.themes.Soft()) as app:
+# ------------ APLICACIÓN GRADIO ------------
+with gr.Blocks(title="MultiModel AI Assistant (SambaNova + Local)", theme=gr.themes.Soft()) as app:
+    gr.Markdown("# 🤖 MultiModel AI Assistant")
+    gr.Markdown("**Características:** ✅ SambaNova API ✅ Modelos Locales ✅ Generación de Imágenes")
+    
     with gr.Row():
-        with gr.Column(scale=1):
-            model_dropdown = gr.Dropdown(
-                choices=[
-                    "llama3", "qwen", "deepseek",
-                    "starcoder2", "deepseek-coder", "bria"
-                ],
-                value="llama3",
-                label="Seleccionar Modelo"
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                height=500,
+                placeholder="Envía un mensaje para comenzar...",
+                show_copy_button=True
+            )
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    label="Tu mensaje",
+                    placeholder="Escribe tu pregunta aquí...",
+                    scale=4
+                )
+                submit_btn = gr.Button("Enviar 🚀", variant="primary", scale=1)
+            
+            img_input = gr.Image(
+                type="pil", 
+                label="Subir imagen (opcional)",
+                height=100
             )
-            gr.Markdown("### Modelos Disponibles")
-            gr.Markdown("""
-            - Chat: Llama3, Qwen, DeepSeek  
-            - Código: StarCoder2, DeepSeek-Coder  
-            - Imágenes: BRIA
-            """)
-
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=500)
-            msg = gr.Textbox(label="Mensaje")
-            img_input = gr.Image(type="pil", label="Subir imagen (opcional)")
-            submit_btn = gr.Button("Enviar", variant="primary")
+    
+    # Estado del sistema
+    status = "✅ SambaNova: " + ("Conectado" if tools.sn_client else "No disponible")
+    gr.Markdown(f"**Estado del sistema:** {status}")
 
     submit_btn.click(
         process_input,
-        inputs=[msg, img_input, model_dropdown, chatbot],
+        inputs=[msg, img_input, chatbot],
+        outputs=chatbot
+    )
+    
+    # Enter para enviar
+    msg.submit(
+        process_input,
+        inputs=[msg, img_input, chatbot],
         outputs=chatbot
     )
 
 if __name__ == "__main__":
-    app.launch(server_port=7860)
-    
+    print("🚀 Iniciando aplicación...")
+    app.launch(server_port=7860, share=False)