Update app.py

app.py

@@ -1,19 +1,27 @@
 import os
-import torch
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-
-# 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, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    trust_remote_code=True
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+
+REPO_ID = "unsloth/Kimi-K2.6-GGUF"
+# Nota: Si el modelo está shardeado en varias partes (-00001-of-00014),
+# llama.cpp requiere que descargues la primera parte y automáticamente detecta el resto,
+# o debes usar un archivo único cuantizado menor (p. ej. Q2_K o un modelo de menor tamaño).
+FILENAME = "UD-Q4_K_XL/Kimi-K2.6-UD-Q4_K_XL-00001-of-00014.gguf"
+
+# Descargar archivo del Hub
+model_path = hf_hub_download(
+    repo_id=REPO_ID,
+    filename=FILENAME
+)
+
+# Inicializar motor GGUF con soporte GPU
+llm = Llama(
+    model_path=model_path,
+    n_gpu_layers=-1,       # Enviar capas a la GPU
+    n_ctx=4096,
+    verbose=False
 )
 
 @spaces.GPU(duration=120)
@@ -28,34 +36,27 @@ def generate_response(message, history, system_prompt=""):
             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)
-    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,
+
+    response_stream = llm.create_chat_completion(
+        messages=messages,
+        max_tokens=2048,
         temperature=0.6,
         top_p=0.9,
+        stream=True
     )
-    
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
-    
+
     partial_text = ""
-    for new_token in streamer:
-        partial_text += new_token
-        yield partial_text
+    for chunk in response_stream:
+        delta = chunk["choices"][0]["delta"]
+        if "content" in delta:
+            partial_text += delta["content"]
+            yield partial_text
 
-# Interfaz Gradio
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Kimi-K2.6 ZeroGPU Service",
+    title="Kimi-K2.6 Inference",
     additional_inputs=[
-        gr.Textbox("Eres un asistente experto en programación y resolución de tareas.", label="System Prompt")
+        gr.Textbox("Eres un asistente experto en programación.", label="System Prompt")
     ]
 )