Update app.py

app.py

@@ -2,7 +2,7 @@
 import os, json, torch, asyncio
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, StaticCache # Added StaticCache
 from snac import SNAC
 
 # 0) Login + Device ---------------------------------------------------
@@ -100,22 +100,28 @@ async def tts(ws: WebSocket):
         voice = req.get("voice", "Jakob")
 
         ids, attn  = build_prompt(text, voice)
+        # Initialized StaticCache
+        past = StaticCache(config=model.config, max_batch_size=1, max_cache_len=ids.size(1) + CHUNK_TOKENS, device=device, dtype=model.dtype)
         offset_len = ids.size(1)          # wie viele Tokens existieren schon
+        last_tok   = None # Initialized last_tok
         buf         = []
 
         while True:
-            # --- Mini‑Generate (Cache Disabled) -------------------------------------------
+            print(f"DEBUG: Before generate - past is None: {past is None}", flush=True) # Added logging
+            print(f"DEBUG: Before generate - type of past: {type(past) if past is not None else 'None'}", flush=True) # Added logging
+            # --- Mini‑Generate (Cache Re-enabled) -------------------------------------------
             gen = model.generate(
-                input_ids       = ids,
-                attention_mask  = attn,
-                past_key_values = None, # Cache disabled
-                max_new_tokens = 1,
+                input_ids       = ids if past.get_seq_length() == 0 else torch.tensor([[last_tok]], device=device), # Use cache seq length
+                attention_mask  = attn if past.get_seq_length() == 0 else None, # Use cache seq length
+                past_key_values = past, # Use StaticCache instance
+                max_new_tokens = 1, # Set max_new_tokens to 1 for debugging cache
                 logits_processor=[masker],
-                do_sample=False, temperature=0.7, top_p=0.95,
-                use_cache=False, # Cache disabled
+                do_sample=True, temperature=0.7, top_p=0.95,
+                use_cache=True, # Re-enabled cache
                 return_dict_in_generate=True,
                 return_legacy_cache=True
             )
+            print(f"DEBUG: After generate - type of gen.past_key_values: {type(gen.past_key_values)}", flush=True) # Added logging
 
             # ----- neue Tokens heraus schneiden --------------------------
             seq  = gen.sequences[0].tolist()