Update app.py

app.py

@@ -12,8 +12,7 @@ import threading
 logging.basicConfig(level=logging.INFO)
 
 # --- MODEL MAP ---
-# This maps the "light", "medium", "heavy" keys from your frontend
-# to the actual model files on Hugging Face.
+# Maps the frontend keys to the real model files
 MODEL_MAP = {
     "light": {
         "repo_id": "TheBloke/stablelm-zephyr-3b-GGUF",
@@ -31,7 +30,7 @@ MODEL_MAP = {
 
 # --- GLOBAL CACHE & LOCK ---
 llm_cache = {} # Caches loaded models
-model_lock = threading.Lock() # Prevents two requests from using the model at once
+model_lock = threading.Lock() # Prevents crashes
 
 app = FastAPI()
 
@@ -46,11 +45,6 @@ app.add_middleware(
 
 # --- Helper Function to Load Model ---
 def get_llm_instance(choice: str) -> Llama:
-    """
-    Loads a model based on the choice.
-    Uses hf_hub_download.
-    Caches the loaded model in memory.
-    """
     if choice not in MODEL_MAP:
         logging.error(f"Invalid model choice: {choice}")
         return None
@@ -116,14 +110,13 @@ def get_status():
     return {
         "status": "AI server is online",
         "model_loaded": loaded_model,
-        "models": list(MODEL_MAP.keys()) # <-- This is the CRUCIAL line for your frontend
+        "models": list(MODEL_MAP.keys()) # This is the CRUCIAL line for your frontend
     }
 
 @app.post("/generate")
 async def generate_story(prompt: StoryPrompt):
     """
     Main generation endpoint.
-    Uses the thread lock to ensure stability.
     """
     logging.info("Request received. Waiting to acquire model lock...")
     with model_lock:
@@ -134,7 +127,6 @@ async def generate_story(prompt: StoryPrompt):
                 logging.error(f"Failed to get model for choice: {prompt.model_choice}")
                 return JSONResponse(status_code=503, content={"error": "The AI model is not available or failed to load."})
 
-            # Format the prompt (Zephyr/ChatML format)
             final_prompt = f"""<|user|>
 Here is the story so far:
 {prompt.story_memory}
@@ -165,4 +157,4 @@ Generate the next part of the story.<|endoftext|>
             logging.error(f"An internal error occurred during generation: {e}", exc_info=True)
             return JSONResponse(status_code=500, content={"error": "An unexpected error occurred."})
         finally:
-            logging.info("Releasing model lock.")
+            logging.info("Releasing model lock.")