reverse fix

acestep/constrained_logits_processor.py

@@ -20,12 +20,6 @@ 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
 # ==============================================================================
@@ -520,34 +514,21 @@ 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+)\|>$')
-        out_of_range_count = 0
+        audio_code_pattern = re.compile(r'^<\|audio_code_\d+\|>$')
         
         # Iterate through vocabulary to find audio code tokens
         for token_id in range(self.vocab_size):
             try:
                 token_text = self.tokenizer.decode([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})")
+                if audio_code_pattern.match(token_text):
+                    self.audio_code_token_ids.add(token_id)
             except Exception:
                 continue
         
         if self.debug:
-            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}]")
+            logger.debug(f"Found {len(self.audio_code_token_ids)} audio code tokens")
     
     def _build_audio_code_mask(self):
         """
@@ -1522,24 +1503,6 @@ 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,32 +774,11 @@ 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|>.
-        Clamps code values to valid range [0, 63999].
-        """
+        """Extract integer audio codes from prompt tokens like <|audio_code_123|>."""
         if not code_str:
             return []
         try:
-            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
+            return [int(x) for x in re.findall(r"<\|audio_code_(\d+)\|>", code_str)]
         except Exception as e:
             logger.debug(f"[_parse_audio_code_string] Failed to parse audio code string: {e}")
             return []
@@ -821,23 +800,16 @@ class AceStepHandler:
                 detokenizer = self.model.detokenizer
                 
                 # Get codebook size for validation
-                # DIT quantizer supports codebook size = 64000 (valid range: 0-63999)
-                MAX_AUDIO_CODE = 63999
-                codebook_size = getattr(quantizer, 'codebook_size', 64000)
+                codebook_size = getattr(quantizer, 'codebook_size', 65536)
                 if hasattr(quantizer, 'quantizers') and len(quantizer.quantizers) > 0:
                     codebook_size = getattr(quantizer.quantizers[0], 'codebook_size', 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]
+                # Validate code IDs are within valid range
+                invalid_codes = [c for c in code_ids if c < 0 or c >= codebook_size]
                 if invalid_codes:
-                    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]
+                    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]
                 
                 num_quantizers = getattr(quantizer, "num_quantizers", 1)
                 # Create indices tensor: [T_5Hz]