Update app.py

app.py

@@ -3,6 +3,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.generation.utils import Cache # Added import
 from snac import SNAC
 
 # 0) Login + Device ---------------------------------------------------
@@ -109,17 +110,33 @@ async def tts(ws: WebSocket):
             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 if past is None else torch.tensor([[last_tok]], device=device), # Re-enabled cache input
-                attention_mask  = attn if past is None else None, # Re-enabled cache attention
-                past_key_values = past, # Re-enabled cache
-                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
-            )
+            # --- Mini‑Generate (Cache Re-enabled) -------------------------------------------
+            if past is None:
+                gen = model.generate(
+                    input_ids       = ids,
+                    attention_mask  = attn,
+                    past_key_values = past,
+                    max_new_tokens = 1,
+                    logits_processor=[masker],
+                    do_sample=True, temperature=0.7, top_p=0.95,
+                    use_cache=True,
+                    return_dict_in_generate=True,
+                )
+            else:
+                # Provide attention mask for the single new token
+                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,
+                    max_new_tokens = 1,
+                    logits_processor=[masker],
+                    do_sample=True, temperature=0.7, top_p=0.95,
+                    use_cache=True,
+                    return_dict_in_generate=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 --------------------------