Update app.py

app.py

@@ -4,6 +4,7 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
 from transformers.generation.utils import Cache # Added import
+from transformers.cache_utils import DynamicCache # Added import
 from snac import SNAC
 
 # 0) Login + Device ---------------------------------------------------
@@ -108,14 +109,14 @@ async def tts(ws: WebSocket):
 
         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
-            # --- Mini‑Generate (Cache Re-enabled) -------------------------------------------
-            # --- Mini‑Generate (Cache Re-enabled) -------------------------------------------
+            print(f"DEBUG: Before generate - type of past: {type(past) if past is not None else 'None'}", flush=True) # Added logging)
+
             if past is None:
+                # First generation step
                 gen = model.generate(
                     input_ids       = ids,
                     attention_mask  = attn,
-                    past_key_values = past,
+                    past_key_values = past, # This will be None
                     max_new_tokens = 1,
                     logits_processor=[masker],
                     do_sample=True, temperature=0.7, top_p=0.95,
@@ -123,14 +124,13 @@ async def tts(ws: WebSocket):
                     return_dict_in_generate=True,
                 )
             else:
-                # Provide attention mask for the single new token
+                # Subsequent generation steps
                 current_input_ids = torch.tensor([[last_tok]], device=device)
                 current_attention_mask = torch.ones_like(current_input_ids)
-
                 gen = model.generate(
                     input_ids       = current_input_ids,
                     attention_mask  = current_attention_mask,
-                    past_key_values = past,
+                    past_key_values = past, # This will be a Cache object
                     max_new_tokens = 1,
                     logits_processor=[masker],
                     do_sample=True, temperature=0.7, top_p=0.95,
@@ -138,7 +138,17 @@ async def tts(ws: WebSocket):
                     return_dict_in_generate=True,
                     cache_position=torch.tensor([offset_len], device=device) # Explicitly pass cache_position
                 )
-            print(f"DEBUG: After generate - type of gen.past_key_values: {type(gen.past_key_values)}", flush=True) # Added logging
+
+            print(f"DEBUG: After generate - type of gen.past_key_values: {type(gen.past_key_values)}", flush=True) # Added logging)
+
+            # Convert legacy tuple cache to DynamicCache if necessary (only after the first step)
+            if past is None and isinstance(gen.past_key_values, tuple):
+                past = DynamicCache.from_legacy_cache(gen.past_key_values)
+            else:
+                # For subsequent steps, just update past with the new cache object
+                past = gen.past_key_values
+
+            print(f"DEBUG: After cache update - type of past: {type(past)}", flush=True) # Added logging)
 
             # ----- neue Tokens heraus schneiden --------------------------
             seq  = gen.sequences[0].tolist()
@@ -147,11 +157,7 @@ async def tts(ws: WebSocket):
                 break
             offset_len += len(new)
 
-            # ----- 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 # Re-enabled cache update
-            print(f"DEBUG: After cache update - type of past: {type(past)}", flush=True) # Added logging
+            # ----- Update last_tok ---------
             last_tok = new[-1]
 
             print("new tokens:", new[:25], flush=True)