Update app.py

app.py

@@ -33,11 +33,22 @@ class AudioMask(LogitsProcessor):
         ])
         self.eos   = torch.tensor([EOS_TOKEN], device=audio_ids.device)
         self.sent_blocks = 0
+        self.buffer_pos = 0 # Added buffer position
 
     def __call__(self, input_ids, logits):
-        allowed = self.allow
+        # Calculate allowed tokens based on buffer position
+        start_token = AUDIO_BASE + self.buffer_pos * 4096
+        end_token = start_token + 4096
+        allowed_audio = torch.arange(start_token, end_token, device=self.allow.device)
+
+        allowed = torch.cat([
+            torch.tensor([NEW_BLOCK], device=self.allow.device),
+            allowed_audio
+        ])
+
         if self.sent_blocks:                        # ab 1. Block EOS zulassen
             allowed = torch.cat([allowed, self.eos])
+
         mask = logits.new_full(logits.shape, float("-inf"))
         mask[:, allowed] = 0
         return logits + mask
@@ -105,6 +116,7 @@ async def tts(ws: WebSocket):
         offset_len = ids.size(1)          # wie viele Tokens existieren schon
         last_tok   = None
         buf        = []
+        masker.buffer_pos = 0 # Initialize buffer position for masker
 
         while True:
             # --- Mini‑Generate (Cache Disabled for Debugging) -------------------------------------------
@@ -143,10 +155,12 @@ async def tts(ws: WebSocket):
                     buf.clear()
                     continue
                 buf.append(t - AUDIO_BASE)
+                masker.buffer_pos += 1 # Increment buffer position
                 if len(buf) == 7:
                     await ws.send_bytes(decode_block(buf))
                     buf.clear()
                     masker.sent_blocks = 1      # ab jetzt EOS zulässig
+                    masker.buffer_pos = 0 # Reset buffer position after sending a block
 
     except (StopIteration, WebSocketDisconnect):
         pass