Update app.py

app.py

@@ -1,44 +1,76 @@
 import gradio as gr
 import spaces
 import torch
-import os
+import gc
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from groq import Groq
 
-# --- Configurações Locais (H200) ---
-LOCAL_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
+# --- CATÁLOGO DE MODELOS ---
+# Adicione quantos quiser aqui (que caibam na VRAM um por vez)
+MODEL_MAP = {
+    "qwen-32b": "Qwen/Qwen2.5-Coder-32B-Instruct",
+    "llama-8b": "meta-llama/Llama-3.1-8B-Instruct",
+    "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.3"
+}
+
+# --- Estado Global ---
+current_model_id = None
 model = None
 tokenizer = None
 
-# --- Configuração Groq ---
-# Ele tenta pegar a chave dos segredos do Space
-groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-
-# --- Função 1: Roda na H200 (Gasta Cota) ---
-# Diminuí para 60s para ajudar no seu reset do Colab
-@spaces.GPU(duration=60)
-def run_local_qwen(messages):
+# --- Função de Limpeza de VRAM ---
+def free_memory():
     global model, tokenizer
+    if model is not None:
+        del model
+        del tokenizer
+        gc.collect()
+        torch.cuda.empty_cache()
+        print("🧹 VRAM limpa!")
+
+# --- A Mágica do Roteamento na GPU ---
+# Aumentei a duration para 90s porque trocar de modelo gasta uns 20s
+@spaces.GPU(duration=90)
+def router(message, history, model_name_key):
+    global model, tokenizer, current_model_id
+
+    target_id = MODEL_MAP.get(model_name_key)
     
-    # Lazy Loading
-    if model is None:
-        print(f"🚀 Carregando {LOCAL_MODEL_ID} na H200...")
-        tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_ID)
+    if not target_id:
+        return f"❌ Erro: Modelo '{model_name_key}' não encontrado no catálogo."
+
+    # --- LÓGICA DE SWAP (TROCA) ---
+    if current_model_id != target_id:
+        print(f"🔄 Trocando de {current_model_id} para {target_id}...")
+        free_memory() # Esvazia a GPU
+        
+        print("🚀 Carregando novo modelo...")
+        tokenizer = AutoTokenizer.from_pretrained(target_id)
         model = AutoModelForCausalLM.from_pretrained(
-            LOCAL_MODEL_ID,
+            target_id,
             torch_dtype=torch.bfloat16,
             device_map="cuda"
         )
-    
-    # Prepara prompt
+        current_model_id = target_id
+        print("✅ Modelo carregado!")
+    else:
+        print("⚡ Modelo já está na VRAM. Usando cache.")
+
+    # --- INFERÊNCIA ---
+    # Formata histórico
+    messages = []
+    for user_msg, bot_msg in history:
+        if user_msg: messages.append({"role": "user", "content": user_msg})
+        if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
     text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
+    
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Gera
     outputs = model.generate(
         **inputs, 
         max_new_tokens=1024,
@@ -49,64 +81,21 @@ def run_local_qwen(messages):
     response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return response
 
-# --- Função 2: Roda no Groq (NÃO Gasta Cota da GPU) ---
-def run_groq(messages, model_id="llama3-70b-8192"):
-    print(f"⚡ Chamando Groq: {model_id}...")
-    try:
-        completion = groq_client.chat.completions.create(
-            model=model_id,
-            messages=messages,
-            temperature=0.7,
-            max_tokens=1024,
-            top_p=1,
-            stream=False,
-            stop=None,
-        )
-        return completion.choices[0].message.content
-    except Exception as e:
-        return f"❌ Erro no Groq: {str(e)}"
-
-# --- O Roteador Central (A Inteligência) ---
-def router(message, history, model_selector):
-    # Formata histórico para padrão OpenAI/Groq
-    messages = []
-    for user_msg, bot_msg in history:
-        if user_msg: messages.append({"role": "user", "content": user_msg})
-        if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
-    messages.append({"role": "user", "content": message})
-
-    # A Lógica de Roteamento
-    if model_selector == "Local: Qwen 2.5 32B (H200)":
-        return run_local_qwen(messages)
-    
-    elif model_selector == "Groq: Llama 3 70B":
-        return run_groq(messages, "llama3-70b-8192")
-        
-    elif model_selector == "Groq: Mixtral 8x7B":
-        return run_groq(messages, "mixtral-8x7b-32768")
-    
-    else:
-        return "Modelo não reconhecido."
-
 # --- Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔀 APIDOST Router")
-    gr.Markdown("Roteamento híbrido: H200 Local (ZeroGPU) + Groq Cloud (LPU)")
+    gr.Markdown("# 🧠 Gabriel's Multi-Model Switcher")
     
     with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=[
-                "Local: Qwen 2.5 32B (H200)", 
-                "Groq: Llama 3 70B", 
-                "Groq: Mixtral 8x7B"
-            ],
-            value="Groq: Llama 3 70B", # Padrão no Groq pra economizar sua cota
-            label="Escolha o Cérebro"
+        # Dropdown para escolher qual modelo do HF carregar
+        model_selector = gr.Dropdown(
+            choices=list(MODEL_MAP.keys()),
+            value="qwen-32b",
+            label="Escolha o Modelo (Isso faz swap na GPU)"
         )
     
     chat = gr.ChatInterface(
         fn=router,
-        additional_inputs=[model_dropdown] # Passa o dropdown pro router
+        additional_inputs=[model_selector]
     )
 
 if __name__ == "__main__":