>w<

app.py

@@ -5,7 +5,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
 
-MODEL_ID = "OpceanAI/Yuuki-best"
+# Definir todos los modelos disponibles
+MODELS = {
+    "yuuki-best": "OpceanAI/Yuuki-best",
+    "yuuki-3.7": "OpceanAI/Yuuki-3.7",
+    "yuuki-v0.1": "OpceanAI/Yuuki-v0.1"
+}
 
 app = FastAPI(
     title="Yuuki API",
@@ -20,21 +25,31 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-print(f"Loading tokenizer from {MODEL_ID}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# Cache de modelos cargados
+loaded_models = {}
+loaded_tokenizers = {}
 
-print(f"Loading model from {MODEL_ID}...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float32
-).to("cpu")
 
-model.eval()
-print("Model ready!")
+def load_model(model_key: str):
+    """Lazy load: solo carga el modelo cuando se necesita"""
+    if model_key not in loaded_models:
+        print(f"Loading {model_key}...")
+        model_id = MODELS[model_key]
+        
+        loaded_tokenizers[model_key] = AutoTokenizer.from_pretrained(model_id)
+        loaded_models[model_key] = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float32
+        ).to("cpu")
+        loaded_models[model_key].eval()
+        print(f"{model_key} ready!")
+    
+    return loaded_models[model_key], loaded_tokenizers[model_key]
 
 
 class GenerateRequest(BaseModel):
     prompt: str = Field(..., min_length=1, max_length=4000)
+    model: str = Field(default="yuuki-best", description="Model to use")
     max_new_tokens: int = Field(default=120, ge=1, le=512)
     temperature: float = Field(default=0.7, ge=0.1, le=2.0)
     top_p: float = Field(default=0.95, ge=0.0, le=1.0)
@@ -42,6 +57,7 @@ class GenerateRequest(BaseModel):
 
 class GenerateResponse(BaseModel):
     response: str
+    model: str
     tokens_generated: int
     time_ms: int
 
@@ -50,9 +66,10 @@ class GenerateResponse(BaseModel):
 def root():
     return {
         "message": "Yuuki Local Inference API",
-        "model": MODEL_ID,
+        "models": list(MODELS.keys()),
         "endpoints": {
             "health": "GET /health",
+            "models": "GET /models",
             "generate": "POST /generate",
             "docs": "GET /docs"
         }
@@ -61,14 +78,38 @@ def root():
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "model": MODEL_ID}
+    return {
+        "status": "ok",
+        "available_models": list(MODELS.keys()),
+        "loaded_models": list(loaded_models.keys())
+    }
+
+
+@app.get("/models")
+def list_models():
+    return {
+        "models": [
+            {"id": key, "name": value} 
+            for key, value in MODELS.items()
+        ]
+    }
 
 
 @app.post("/generate", response_model=GenerateResponse)
 def generate(req: GenerateRequest):
+    # Validar que el modelo existe
+    if req.model not in MODELS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid model. Available: {list(MODELS.keys())}"
+        )
+    
     try:
         start = time.time()
 
+        # Cargar modelo (lazy load)
+        model, tokenizer = load_model(req.model)
+
         inputs = tokenizer(
             req.prompt,
             return_tensors="pt",
@@ -96,6 +137,7 @@ def generate(req: GenerateRequest):
 
         return GenerateResponse(
             response=response_text.strip(),
+            model=req.model,
             tokens_generated=len(new_tokens),
             time_ms=elapsed_ms
         )