Update app.py

app.py

@@ -1,57 +1,63 @@
+import os
+import torch
 import spaces
 import gradio as gr
-import torch
-from fastapi import FastAPI, Request
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
-MODEL_ID = "Qwen/Qwen3.8-27B"  # Modelos >14B suelen exceder la memoria dinámica de ZeroGPU
+# Si usas la versión cuantizada para ZeroGPU:
+MODEL_ID = "unsloth/Kimi-K2.6-GGUF/UD-Q4_K_XL/Kimi-K2.6-UD-Q4_K_XL-00001-of-00014.gguf"  # O un checkpoint cuantizado compatible
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
-    device_map="auto"
+    device_map="auto",
+    trust_remote_code=True
 )
 
-app = FastAPI()
-
 @spaces.GPU(duration=120)
-def generate_response(prompt: str, max_tokens: int = 2048):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_tokens)
+def generate_response(message, history, system_prompt=""):
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
     
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
+            
+    messages.append({"role": "user", "content": message})
     
-    output_text = ""
-    for new_text in streamer:
-        output_text += new_text
-    return output_text
-
-@app.post("/v1/chat/completions")
-async def chat_completions(request: Request):
-    data = await request.json()
-    messages = data.get("messages", [])
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=2048,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9,
+    )
     
-    response_text = generate_response(prompt)
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
     
-    return {
-        "id": "chatcmpl-zerogpu",
-        "object": "chat.completion",
-        "choices": [{
-            "index": 0,
-            "message": {
-                "role": "assistant",
-                "content": response_text
-            },
-            "finish_reason": "stop"
-        }]
-    }
+    partial_text = ""
+    for new_token in streamer:
+        partial_text += new_token
+        yield partial_text
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Qwen ZeroGPU Endpoint")
+# Interfaz Gradio
+demo = gr.ChatInterface(
+    fn=generate_response,
+    title="Kimi-K2.6 ZeroGPU Service",
+    additional_inputs=[
+        gr.Textbox("Eres un asistente experto en programación y resolución de tareas.", label="System Prompt")
+    ]
+)
 
-app = gr.mount_gradio_app(app, demo, path="/")
+if __name__ == "__main__":
+    demo.queue().launch()