fix max audio code id

acestep/constrained_logits_processor.py

@@ -19,6 +19,9 @@ from acestep.constants import (
     VALID_TIME_SIGNATURES,
 )
 
+# Maximum valid audio code value (codebook size = 64000)
+MAX_AUDIO_CODE = 63999
+
 
 # ==============================================================================
 # FSM States for Constrained Decoding
@@ -514,21 +517,61 @@ class MetadataConstrainedLogitsProcessor(LogitsProcessor):
         """
         Precompute audio code token IDs (tokens matching <|audio_code_\\d+|>).
         These tokens should be blocked during caption generation.
+        Only tokens with code values in range [0, MAX_AUDIO_CODE] are included.
         """
         import re
-        audio_code_pattern = re.compile(r'^<\|audio_code_\d+\|>$')
+        audio_code_pattern = re.compile(r'^<\|audio_code_(\d+)\|>$')
+        invalid_tokens_count = 0
         
         # Iterate through vocabulary to find audio code tokens
         for token_id in range(self.vocab_size):
             try:
                 token_text = self.tokenizer.decode([token_id])
-                if audio_code_pattern.match(token_text):
-                    self.audio_code_token_ids.add(token_id)
+                match = audio_code_pattern.match(token_text)
+                if match:
+                    # Extract code value from token text
+                    code_value = int(match.group(1))
+                    # Only add tokens with valid code values (0-63999)
+                    if 0 <= code_value <= MAX_AUDIO_CODE:
+                        self.audio_code_token_ids.add(token_id)
+                    else:
+                        invalid_tokens_count += 1
+                        if self.debug:
+                            logger.debug(f"Skipping audio code token {token_id} with invalid code value {code_value} (max: {MAX_AUDIO_CODE})")
             except Exception:
                 continue
         
-        if self.debug:
-            logger.debug(f"Found {len(self.audio_code_token_ids)} audio code tokens")
+        if invalid_tokens_count > 0:
+            logger.warning(f"Found {invalid_tokens_count} audio code tokens with values outside valid range [0, {MAX_AUDIO_CODE}]")
+        
+        # Log warning if no valid tokens found (this would prevent code generation)
+        if len(self.audio_code_token_ids) == 0:
+            logger.warning(f"No valid audio code tokens found in vocabulary (range [0, {MAX_AUDIO_CODE}]). Code generation may fail.")
+        elif self.debug:
+            logger.debug(f"Found {len(self.audio_code_token_ids)} valid audio code tokens (range [0, {MAX_AUDIO_CODE}])")
+    
+    def _extract_code_from_token(self, token_id: int) -> Optional[int]:
+        """
+        Extract audio code value from a token ID.
+        
+        Args:
+            token_id: Token ID to extract code value from
+            
+        Returns:
+            Code value if token is a valid audio code token, None otherwise
+        """
+        import re
+        audio_code_pattern = re.compile(r'^<\|audio_code_(\d+)\|>$')
+        
+        try:
+            token_text = self.tokenizer.decode([token_id])
+            match = audio_code_pattern.match(token_text)
+            if match:
+                return int(match.group(1))
+        except Exception:
+            pass
+        
+        return None
     
     def _build_audio_code_mask(self):
         """
@@ -1497,6 +1540,8 @@ class MetadataConstrainedLogitsProcessor(LogitsProcessor):
         
         if self.state == FSMState.CODES_GENERATION:
             # Block all non-audio-code tokens (only allow audio codes and EOS)
+            # Note: audio_code_token_ids already contains only valid tokens (0-63999 range)
+            # because _precompute_audio_code_tokens() filters out invalid tokens during initialization
             if self.non_audio_code_mask is not None:
                 # Move mask to same device/dtype as scores if needed
                 if self.non_audio_code_mask.device != scores.device or self.non_audio_code_mask.dtype != scores.dtype:

acestep/handler.py

@@ -774,11 +774,29 @@ class AceStepHandler:
             return None
     
     def _parse_audio_code_string(self, code_str: str) -> List[int]:
-        """Extract integer audio codes from prompt tokens like <|audio_code_123|>."""
+        """Extract integer audio codes from prompt tokens like <|audio_code_123|>.
+        Codes are clamped to valid range [0, 63999] (codebook size = 64000).
+        """
         if not code_str:
             return []
         try:
-            return [int(x) for x in re.findall(r"<\|audio_code_(\d+)\|>", code_str)]
+            codes = [int(x) for x in re.findall(r"<\|audio_code_(\d+)\|>", code_str)]
+            # Clamp codes to valid range [0, 63999]
+            MAX_AUDIO_CODE = 63999
+            clamped_codes = []
+            invalid_codes = []
+            for code in codes:
+                if code < 0 or code > MAX_AUDIO_CODE:
+                    invalid_codes.append(code)
+                    clamped_code = max(0, min(code, MAX_AUDIO_CODE))
+                    clamped_codes.append(clamped_code)
+                else:
+                    clamped_codes.append(code)
+            
+            if invalid_codes:
+                logger.warning(f"[_parse_audio_code_string] Found {len(invalid_codes)} codes outside valid range [0, {MAX_AUDIO_CODE}]: {invalid_codes[:5]}... (clamped to valid range)")
+            
+            return clamped_codes
         except Exception as e:
             logger.debug(f"[_parse_audio_code_string] Failed to parse audio code string: {e}")
             return []
@@ -800,7 +818,8 @@ class AceStepHandler:
                 detokenizer = self.model.detokenizer
                 
                 # Get codebook size for validation
-                codebook_size = getattr(quantizer, 'codebook_size', 65536)
+                # Default to 64000 (codebook size = 64000, valid range = 0-63999)
+                codebook_size = getattr(quantizer, 'codebook_size', 64000)
                 if hasattr(quantizer, 'quantizers') and len(quantizer.quantizers) > 0:
                     codebook_size = getattr(quantizer.quantizers[0], 'codebook_size', codebook_size)