fix max audio code id

acestep/constrained_logits_processor.py

@@ -20,6 +20,12 @@ from acestep.constants import (
 )
 
 
+# ==============================================================================
+# Constants
+# ==============================================================================
+# Maximum valid audio code value (codebook size = 64000, valid range: 0-63999)
+MAX_AUDIO_CODE = 63999
+
 # ==============================================================================
 # FSM States for Constrained Decoding
 # ==============================================================================
@@ -514,21 +520,34 @@ 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+)\|>$')
+        out_of_range_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:
+                        out_of_range_count += 1
+                        if self.debug:
+                            logger.debug(f"Skipping audio code token with out-of-range value: {token_text} (code={code_value})")
             except Exception:
                 continue
         
         if self.debug:
-            logger.debug(f"Found {len(self.audio_code_token_ids)} audio code tokens")
+            logger.debug(f"Found {len(self.audio_code_token_ids)} valid audio code tokens (skipped {out_of_range_count} out-of-range tokens)")
+        if out_of_range_count > 0:
+            logger.warning(f"Skipped {out_of_range_count} audio code tokens with values outside valid range [0, {MAX_AUDIO_CODE}]")
     
     def _build_audio_code_mask(self):
         """
@@ -1503,6 +1522,24 @@ class MetadataConstrainedLogitsProcessor(LogitsProcessor):
                     self.non_audio_code_mask = self.non_audio_code_mask.to(device=scores.device, dtype=scores.dtype)
                 scores = scores + self.non_audio_code_mask
             
+            # Additional validation: block audio code tokens with out-of-range values
+            # This prevents generation of codes > MAX_AUDIO_CODE even if they exist in vocabulary
+            import re
+            audio_code_pattern = re.compile(r'^<\|audio_code_(\d+)\|>$')
+            for token_id in self.audio_code_token_ids:
+                try:
+                    token_text = self.tokenizer.decode([token_id])
+                    match = audio_code_pattern.match(token_text)
+                    if match:
+                        code_value = int(match.group(1))
+                        # Block tokens with code values outside valid range
+                        if code_value > MAX_AUDIO_CODE:
+                            scores[:, token_id] = float('-inf')
+                            if self.debug:
+                                logger.debug(f"Blocking out-of-range audio code token: {token_text} (code={code_value})")
+                except Exception:
+                    continue
+            
             # Apply duration constraint in codes generation phase
             if self.target_codes is not None and self.eos_token_id is not None:
                 if self.codes_count < self.target_codes:

acestep/handler.py

@@ -774,11 +774,32 @@ 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|>.
+        Clamps code values to valid range [0, 63999].
+        """
         if not code_str:
             return []
         try:
-            return [int(x) for x in re.findall(r"<\|audio_code_(\d+)\|>", code_str)]
+            MAX_AUDIO_CODE = 63999
+            codes = [int(x) for x in re.findall(r"<\|audio_code_(\d+)\|>", code_str)]
+            # Clamp codes to valid range [0, 63999]
+            clamped_codes = []
+            clamped_count = 0
+            for code in codes:
+                if code < 0:
+                    clamped_codes.append(0)
+                    clamped_count += 1
+                elif code > MAX_AUDIO_CODE:
+                    clamped_codes.append(MAX_AUDIO_CODE)
+                    clamped_count += 1
+                else:
+                    clamped_codes.append(code)
+            
+            if clamped_count > 0:
+                logger.warning(f"[_parse_audio_code_string] Clamped {clamped_count} audio code values to valid range [0, {MAX_AUDIO_CODE}]")
+            
+            return clamped_codes
         except Exception as e:
             logger.debug(f"[_parse_audio_code_string] Failed to parse audio code string: {e}")
             return []
@@ -800,16 +821,23 @@ class AceStepHandler:
                 detokenizer = self.model.detokenizer
                 
                 # Get codebook size for validation
-                codebook_size = getattr(quantizer, 'codebook_size', 65536)
+                # DIT quantizer supports codebook size = 64000 (valid range: 0-63999)
+                MAX_AUDIO_CODE = 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)
                 
-                # Validate code IDs are within valid range
-                invalid_codes = [c for c in code_ids if c < 0 or c >= codebook_size]
+                # Use 64000 as hard limit regardless of what quantizer reports
+                # This ensures compatibility with the actual DIT quantizer codebook size
+                effective_codebook_size = 64000
+                effective_max_code = MAX_AUDIO_CODE
+                
+                # Validate code IDs are within valid range [0, 63999]
+                invalid_codes = [c for c in code_ids if c < 0 or c > effective_max_code]
                 if invalid_codes:
-                    logger.warning(f"[_decode_audio_codes_to_latents] Found {len(invalid_codes)} invalid codes out of range [0, {codebook_size}): {invalid_codes[:5]}...")
-                    # Clamp invalid codes to valid range
-                    code_ids = [max(0, min(c, codebook_size - 1)) for c in code_ids]
+                    logger.warning(f"[_decode_audio_codes_to_latents] Found {len(invalid_codes)} invalid codes out of range [0, {effective_max_code}]: {invalid_codes[:5]}...")
+                    # Clamp invalid codes to valid range [0, 63999]
+                    code_ids = [max(0, min(c, effective_max_code)) for c in code_ids]
                 
                 num_quantizers = getattr(quantizer, "num_quantizers", 1)
                 # Create indices tensor: [T_5Hz]