Update app.py

app.py

@@ -2,8 +2,8 @@
 import os, json, asyncio, torch, logging
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from transformers import (AutoTokenizer, AutoModelForCausalLM,
-                          LogitsProcessor, generation_utils)
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
+from transformers.generation.utils import Cache          
 from snac import SNAC
 
 # ── 0. Auth & Device ────────────────────────────────────────────────
@@ -58,6 +58,7 @@ async def load_models():
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_REPO,
         low_cpu_mem_usage=True,
+        return_legacy_cache=True,
         device_map={"": 0} if device == "cuda" else None,
         torch_dtype=torch.bfloat16 if device == "cuda" else None,
     )