Update app.py

app.py

@@ -4,48 +4,63 @@ import os
 
 app = Flask(__name__)
 
-MODEL_PATH = "./model.gguf"
-MODEL_URL = "https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat.Q4_K_M.gguf"
-
-# Download if missing
-if not os.path.exists(MODEL_PATH):
-    os.system(f"wget -O {MODEL_PATH} {MODEL_URL}")
-
-# Load model (optimized for CPU Spaces)
-llm = Llama(
-    model_path=MODEL_PATH,
-    n_threads=4,
-    n_ctx=2048,
-    use_mlock=False,
-)
-
-@app.route('/')
-def index():
-    return jsonify({
-        "message": "FugthDes Story Generator Active",
-        "model": "TinyLlama GGUF (CPU)"
-    })
-
-@app.route('/generate', methods=['POST'])
+MODEL_URLS = {
+    "light": "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q3_K_S.gguf",
+    "medium": "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q4_K_M.gguf",
+    "heavy": "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q5_0.gguf"
+}
+
+MODEL_PATHS = {
+    k: f"{k}.gguf" for k in MODEL_URLS
+}
+
+current_model = None
+llm = None
+
+def ensure_model(model_choice):
+    global llm, current_model
+    model_path = MODEL_PATHS[model_choice]
+    url = MODEL_URLS[model_choice]
+
+    if not os.path.exists(model_path):
+        print(f"Downloading {model_choice} model...")
+        os.system(f"wget -O {model_path} {url}")
+
+    if current_model != model_choice:
+        print(f"Loading {model_choice} model...")
+        llm = Llama(model_path=model_path, n_ctx=2048, n_threads=4, use_mlock=False)
+        current_model = model_choice
+    return llm
+
+
+@app.route("/status")
+def status():
+    return jsonify({"status": "ok" if llm else "not_loaded", "model": current_model})
+
+
+@app.route("/generate", methods=["POST"])
 def generate():
-    data = request.get_json()
+    data = request.get_json(force=True)
+    model_choice = data.get("model_choice", "light")
     prompt = data.get("prompt", "")
-    feedback = data.get("feedback", "")
     story_memory = data.get("story_memory", "")
+    feedback = data.get("feedback", "")
+
+    llm = ensure_model(model_choice)
 
-    final_prompt = story_memory + "\n\n" + prompt
+    full_prompt = story_memory + "\n\n" + prompt
     if feedback:
-        final_prompt += f"\n\nUser feedback: {feedback}\nContinue or refine story naturally."
+        full_prompt += f"\n\nUser feedback: {feedback}\n"
+
+    result = llm(full_prompt, max_tokens=512, temperature=0.8)
+    text = result["choices"][0]["text"].strip()
+    return jsonify({"response": text})
 
-    print("Prompt received:", final_prompt[:250])
 
-    output = llm(final_prompt, max_tokens=512, temperature=0.8, top_p=0.9)
-    response_text = output["choices"][0]["text"].strip()
+@app.route("/")
+def root():
+    return "StableLM Zephyr GGUF API running!"
 
-    return jsonify({
-        "response": response_text,
-        "tokens_used": output["usage"]["total_tokens"]
-    })
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)