Update app.py

app.py

@@ -2,27 +2,19 @@ import os
 import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
 # --- Configurações ---
-# Vamos começar com um modelo poderoso que cabe na H200 tranquilo
-MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct" 
-# Ou se quiser algo mais leve: "meta-llama/Llama-3.1-8B-Instruct"
+MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
 
-app = FastAPI(title="APIDOST - Gabriel's Router")
+# Variáveis globais iniciadas como None (vazias)
+model = None
+tokenizer = None
 
-print(f"🔄 Carregando modelo: {MODEL_ID}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
-print("✅ Modelo carregado e pronto para a guerra!")
+app = FastAPI(title="APIDOST - Gabriel's Router")
 
-# --- Estruturas de Dados (Schema OpenAI-like) ---
+# --- Estruturas de Dados ---
 class Message(BaseModel):
     role: str
     content: str
@@ -32,13 +24,25 @@ class ChatCompletionRequest(BaseModel):
     messages: list[Message]
     max_tokens: int = 1024
     temperature: float = 0.7
-    stream: bool = False
 
-# --- A Mágica do ZeroGPU ---
-# O decorator @spaces.GPU garante que essa função rode na H200
-@spaces.GPU
+# --- A Mágica do ZeroGPU com Lazy Loading ---
+# duration=120 garante 2 minutos de GPU, tempo suficiente pro load + inferencia
+@spaces.GPU(duration=120) 
 def generate_response(messages, max_tokens, temperature):
-    # Formata o prompt (chat template)
+    global model, tokenizer
+    
+    # O PULO DO GATO: Só carrega se ainda não estiver na memória
+    if model is None:
+        print(f"⏳ Cold Start: Carregando {MODEL_ID} para a VRAM...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.bfloat16,
+            device_map="cuda" # Força o uso da GPU alocada pelo spaces
+        )
+        print("✅ Modelo carregado com sucesso!")
+
+    # Prepara o prompt
     text_prompt = tokenizer.apply_chat_template(
         messages, 
         tokenize=False, 
@@ -47,40 +51,32 @@ def generate_response(messages, max_tokens, temperature):
     
     inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)
 
-    # Configuração de geração
-    generate_kwargs = dict(
-        inputs,
+    # Gera
+    output = model.generate(
+        **inputs,
         max_new_tokens=max_tokens,
         temperature=temperature,
         do_sample=True,
         top_p=0.9,
     )
-
-    # Gera a resposta
-    output = model.generate(**generate_kwargs)
-    response_text = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     
+    response_text = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return response_text
 
 # --- Endpoints ---
-
 @app.get("/")
 def read_root():
-    return {"status": "APIDOST is online", "hardware": "Nvidia H200 (ZeroGPU)"}
+    return {"status": "APIDOST is online", "mode": "Lazy Loading Active"}
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: ChatCompletionRequest):
-    """
-    Endpoint compatível (simplificado) com OpenAI.
-    """
     try:
-        # Converte as mensagens do Pydantic para lista de dicts
+        # Converte mensagens
         msgs = [{"role": m.role, "content": m.content} for m in request.messages]
         
-        # Chama a GPU
+        # Chama a função protegida pelo @spaces.GPU
         response_content = generate_response(msgs, request.max_tokens, request.temperature)
 
-        # Formata a resposta estilo OpenAI
         return {
             "id": "chatcmpl-apidost",
             "object": "chat.completion",
@@ -93,18 +89,13 @@ async def chat_completions(request: ChatCompletionRequest):
                     "content": response_content
                 },
                 "finish_reason": "stop"
-            }],
-            "usage": {
-                "prompt_tokens": 0, # Implementar contagem real se quiser
-                "completion_tokens": 0,
-                "total_tokens": 0
-            }
+            }]
         }
 
     except Exception as e:
+        print(f"❌ Erro: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-# Para rodar localmente ou no Spaces via Docker
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)