Update app.py

app.py

@@ -3,8 +3,6 @@ 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 transformers.cache_utils import DynamicCache # Added import
 from snac import SNAC
 
 # 0) Login + Device ---------------------------------------------------
@@ -102,92 +100,48 @@ async def tts(ws: WebSocket):
         voice = req.get("voice", "Jakob")
 
         ids, attn  = build_prompt(text, voice)
-        past       = None
         offset_len = ids.size(1)          # wie viele Tokens existieren schon
-        last_tok   = None
         buf         = []
 
-        # Initial generation step using model.generate
-        with torch.no_grad():
+        while True:
+            # --- Mini‑Generate (Cache Disabled) -------------------------------------------
             gen = model.generate(
                 input_ids       = ids,
                 attention_mask  = attn,
-                past_key_values = None, # Initial call, no past cache
+                past_key_values = None, # Cache disabled
                 max_new_tokens = 1,
                 logits_processor=[masker],
                 do_sample=True, temperature=0.7, top_p=0.95,
-                use_cache=True,
+                use_cache=False, # Cache disabled
                 return_dict_in_generate=True,
+                return_legacy_cache=True
             )
 
-        # Get the initial cache and last token
-        past = gen.past_key_values
-        if isinstance(past, tuple):
-             past = DynamicCache.from_legacy_cache(past) # Convert legacy tuple cache
-        last_tok = gen.sequences[0].tolist()[-1]
-        offset_len += 1 # Increment offset for the first generated token
-
-        print(f"DEBUG: After initial generate - type of past: {type(past)}", flush=True) # Added logging
-        print("new tokens:", [last_tok], flush=True) # Log the first token
-
-        # Handle the first generated token
-        if last_tok == EOS_TOKEN:
-            raise StopIteration
-        if last_tok == NEW_BLOCK:
-            buf.clear()
-        else:
-            buf.append(last_tok - AUDIO_BASE)
-            if len(buf) == 7:
-                await ws.send_bytes(decode_block(buf))
-                buf.clear()
-                masker.sent_blocks = 1
-
-        # Manual generation loop for subsequent tokens
-        while True:
-            print(f"DEBUG: Before forward - type of past: {type(past)}", flush=True) # Added logging
-
-            # Prepare inputs for the next token
-            current_input_ids = torch.tensor([[last_tok]], device=device)
-            current_attention_mask = torch.ones_like(current_input_ids)
-            current_cache_position = torch.tensor([offset_len], device=device)
-
-            # Perform forward pass
-            with torch.no_grad():
-                 outputs = model(
-                    input_ids=current_input_ids,
-                    attention_mask=current_attention_mask,
-                    past_key_values=past,
-                    cache_position=current_cache_position,
-                    use_cache=True,
-                 )
-
-            # Sample the next token (greedy sampling)
-            next_token_logits = outputs.logits[:, -1, :]
-            # Apply logits processor manually
-            processed_logits = masker(current_input_ids, next_token_logits.unsqueeze(0))[0]
-            next_token_id = torch.argmax(processed_logits).item()
-
-            print(f"DEBUG: After forward - type of outputs.past_key_values: {type(outputs.past_key_values)}", flush=True) # Added logging
-
-            # Update cache and last token
-            past = outputs.past_key_values
-            last_tok = next_token_id
-            offset_len += 1 # Increment offset for the new token
-
-            print("new tokens:", [last_tok], flush=True) # Log the new token
+            # ----- neue Tokens heraus schneiden --------------------------
+            seq  = gen.sequences[0].tolist()
+            new  = seq[offset_len:]
+            if not new:                         # nichts -> fertig
+                break
+            offset_len += len(new)
 
-            # ----- Token‑Handling ----------------------------------------
-            if last_tok == EOS_TOKEN:
-                raise StopIteration
-            if last_tok == NEW_BLOCK:
-                buf.clear()
-                continue # Continue loop to generate the next token
-            buf.append(last_tok - AUDIO_BASE)
-            if len(buf) == 7:
-                await ws.send_bytes(decode_block(buf))
-                buf.clear()
-                masker.sent_blocks = 1      # ab jetzt EOS zulässig
+            # ----- Update ids and attn for next iteration (Cache Disabled) ---------
+            ids = torch.tensor([seq], device=device)
+            attn = torch.ones_like(ids)
 
+            print("new tokens:", new[:25], flush=True)
+
+            # ----- Token‑Handling ----------------------------------------
+            for t in new:
+                if t == EOS_TOKEN:
+                    raise StopIteration
+                if t == NEW_BLOCK:
+                    buf.clear()
+                    continue
+                buf.append(t - AUDIO_BASE)
+                if len(buf) == 7:
+                    await ws.send_bytes(decode_block(buf))
+                    buf.clear()
+                    masker.sent_blocks = 1      # ab jetzt EOS zulässig
 
     except (StopIteration, WebSocketDisconnect):
         pass