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, StaticCache # Added StaticCache
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, DynamicCache # Added StaticCache
 from snac import SNAC
 
 # 0) Login + Device ---------------------------------------------------
@@ -104,7 +104,7 @@ async def tts(ws: WebSocket):
         offset_len = ids.size(1)          # wie viele Tokens existieren schon
         last_tok   = None # Initialized last_tok
         buf         = []
-
+        past_key_values = DynamicCache()
         while True:
             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
@@ -112,14 +112,14 @@ async def tts(ws: WebSocket):
             gen = model.generate(
                 input_ids       = ids if past is None else torch.tensor([[last_tok]], device=device), # Use past is None check
                 attention_mask  = attn if past is None else None, # Use past is None check
-                past_key_values = past, # Pass past (will be None initially, then the cache object)
+                past_key_values = past_key_values, # Pass past (will be None initially, then the cache object)
                 max_new_tokens = 1, # Set max_new_tokens to 1 for debugging cache
                 logits_processor=[masker],
                 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,
-                cache_implementation="static" # Enabled StaticCache via implementation
+                #return_legacy_cache=True,
+                #cache_implementation="static" # Enabled StaticCache via implementation
             )
             print(f"DEBUG: After generate - type of gen.past_key_values: {type(gen.past_key_values)}", flush=True) # Added logging
 
@@ -133,8 +133,12 @@ async def tts(ws: WebSocket):
             # ----- Update past and last_tok (Cache Re-enabled) ---------
             # ids = torch.tensor([seq], device=device) # Removed full sequence update
             # attn = torch.ones_like(ids) # Removed full sequence update
-            past = gen.past_key_values # Update past with the cache object returned by generate
+            #pkv = gen.past_key_values # Update past with the cache object returned by generate
+            print(f"DEBUG: After cache update - type of past: {type(past)}", flush=True) # Added logging
+            #if isinstance(pkv, StaticCache): pkv = pkv.to_legacy()
+            past = gen.past_key_values
             print(f"DEBUG: After cache update - type of past: {type(past)}", flush=True) # Added logging
+
             last_tok = new[-1]
 
             print("new tokens:", new[:25], flush=True)