Delete processing_spark_tts.py

processing_spark_tts.py

@@ -1,889 +0,0 @@
-# coding=utf-8
-# Copyright 2025 SparkAudio & The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Processor class for SparkTTS. Combines text tokenization and audio feature extraction/processing.
-"""
-
-import os # Needed for save_pretrained
-import re # For decoding
-import torch
-import numpy as np
-import soundfile as sf # For audio loading
-import soxr # For resampling
-
-from pathlib import Path
-from typing import Optional, Union, List, Dict, Tuple, Any
-
-from transformers.processing_utils import ProcessorMixin
-from transformers.tokenization_utils_base import BatchEncoding # Return type hint
-from transformers.feature_extraction_utils import BatchFeature # Return type hint
-from transformers.models.auto.tokenization_auto import AutoTokenizer
-from transformers.models.wav2vec2.feature_extraction_wav2vec2 import Wav2Vec2FeatureExtractor
-from transformers.utils import logging, PushToHubMixin # Added PushToHubMixin
-from numpy.lib.stride_tricks import sliding_window_view
-import soxr
-import soundfile
-import random
-
-# Import custom config if needed for defaults
-from .configuration_spark_tts import SparkTTSConfig
-
-logger = logging.get_logger(__name__)
-
-
-# =============================================================================
-# >> START: PASTE CODE FROM sparktts/utils/* HERE <<
-# =============================================================================
-# IMPORTANT: Utility functions needed for processing (audio loading, token parsing)
-# must be defined or imported here.
-
-# --- Paste sparktts/utils/audio.py content here ---
-
-def audio_volume_normalize(audio: np.ndarray, coeff: float = 0.2) -> np.ndarray:
-    """
-    Normalize the volume of an audio signal.
-
-    Parameters:
-        audio (numpy array): Input audio signal array.
-        coeff (float): Target coefficient for normalization, default is 0.2.
-
-    Returns:
-        numpy array: The volume-normalized audio signal.
-    """
-    # Sort the absolute values of the audio signal
-    temp = np.sort(np.abs(audio))
-
-    # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1
-    if temp[-1] < 0.1:
-        scaling_factor = max(
-            temp[-1], 1e-3
-        )  # Prevent division by zero with a small constant
-        audio = audio / scaling_factor * 0.1
-
-    # Filter out values less than 0.01 from temp
-    temp = temp[temp > 0.01]
-    L = temp.shape[0]  # Length of the filtered array
-
-    # If there are fewer than or equal to 10 significant values, return the audio without further processing
-    if L <= 10:
-        return audio
-
-    # Compute the average of the top 10% to 1% of values in temp
-    volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)])
-
-    # Normalize the audio to the target coefficient level, clamping the scale factor between 0.1 and 10
-    audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10)
-
-    # Ensure the maximum absolute value in the audio does not exceed 1
-    max_value = np.max(np.abs(audio))
-    if max_value > 1:
-        audio = audio / max_value
-
-    return audio
-
-
-def load_audio(
-    adfile: Path,
-    sampling_rate: int = None,
-    length: int = None,
-    volume_normalize: bool = False,
-    segment_duration: int = None,
-) -> np.ndarray:
-    r"""Load audio file with target sampling rate and lsength
-
-    Args:
-        adfile (Path): path to audio file.
-        sampling_rate (int, optional): target sampling rate. Defaults to None.
-        length (int, optional): target audio length. Defaults to None.
-        volume_normalize (bool, optional): whether perform volume normalization. Defaults to False.
-        segment_duration (int): random select a segment with duration of {segment_duration}s.
-                                Defualt to None which means the whole audio will be used.
-
-    Returns:
-        audio (np.ndarray): audio
-    """
-
-    audio, sr = soundfile.read(adfile)
-    if len(audio.shape) > 1:
-        audio = audio[:, 0]
-
-    if sampling_rate is not None and sr != sampling_rate:
-        audio = soxr.resample(audio, sr, sampling_rate, quality="VHQ")
-        sr = sampling_rate
-
-    if segment_duration is not None:
-        seg_length = int(sr * segment_duration)
-        audio = random_select_audio_segment(audio, seg_length)
-
-    # Audio volume normalize
-    if volume_normalize:
-        audio = audio_volume_normalize(audio)
-    # check the audio length
-    if length is not None:
-        assert abs(audio.shape[0] - length) < 1000
-        if audio.shape[0] > length:
-            audio = audio[:length]
-        else:
-            audio = np.pad(audio, (0, int(length - audio.shape[0])))
-    return audio
-
-
-def random_select_audio_segment(audio: np.ndarray, length: int) -> np.ndarray:
-    """get an audio segment given the length
-
-    Args:
-        audio (np.ndarray):
-        length (int): audio length = sampling_rate * duration
-    """
-    if audio.shape[0] < length:
-        audio = np.pad(audio, (0, int(length - audio.shape[0])))
-    start_index = random.randint(0, audio.shape[0] - length)
-    end_index = int(start_index + length)
-
-    return audio[start_index:end_index]
-
-def get_ref_clip(wav: np.ndarray, config) -> np.ndarray: # Needs access to config attributes
-    """Get reference audio clip for speaker embedding."""
-    # Make sure config has sample_rate, ref_segment_duration, latent_hop_length
-    if not all(hasattr(config, attr) for attr in ['sample_rate', 'ref_segment_duration', 'latent_hop_length']):
-        raise AttributeError("Config object missing required attributes for get_ref_clip")
-    ref_segment_length = (
-        int(config.sample_rate * config.ref_segment_duration)
-        // config.latent_hop_length
-        * config.latent_hop_length
-    )
-    wav_length = len(wav)
-    if ref_segment_length > wav_length:
-        wav = np.tile(wav, ref_segment_length // wav_length + 1)
-    return wav[:ref_segment_length]
-
-
-# --- Paste sparktts/utils/token_parser.py content here ---
-
-TASK_TOKEN_MAP = {
-    "vc": "<|task_vc|>",
-    "tts": "<|task_tts|>",
-    "asr": "<|task_asr|>",
-    "s2s": "<|task_s2s|>",
-    "t2s": "<|task_t2s|>",
-    "understand": "<|task_understand|>",
-    "caption": "<|task_cap|>",
-    "controllable_tts": "<|task_controllable_tts|>",
-    "prompt_tts": "<|task_prompt_tts|>",
-    "speech_edit": "<|task_edit|>",
-}
-
-LEVELS_MAP = {
-    "very_low": 0,
-    "low": 1,
-    "moderate": 2,
-    "high": 3,
-    "very_high": 4,
-}
-
-LEVELS_MAP_UI = {
-    1: 'very_low',
-    2: 'low',
-    3: 'moderate',
-    4: 'high',
-    5: 'very_high'
-}
-
-GENDER_MAP = {
-    "female": 0,
-    "male": 1,
-}
-
-AGE_MAP = {"Child": 0, "Teenager": 1, "Youth-Adult": 2, "Middle-aged": 3, "Elderly": 4}
-
-EMO_MAP = {
-    "UNKNOWN": 0,
-    "NEUTRAL": 1,
-    "ANGRY": 2,
-    "HAPPY": 3,
-    "SAD": 4,
-    "FEARFUL": 5,
-    "DISGUSTED": 6,
-    "SURPRISED": 7,
-    "SARCASTIC": 8,
-    "EXCITED": 9,
-    "SLEEPY": 10,
-    "CONFUSED": 11,
-    "EMPHASIS": 12,
-    "LAUGHING": 13,
-    "SINGING": 14,
-    "WORRIED": 15,
-    "WHISPER": 16,
-    "ANXIOUS": 17,
-    "NO-AGREEMENT": 18,
-    "APOLOGETIC": 19,
-    "CONCERNED": 20,
-    "ENUNCIATED": 21,
-    "ASSERTIVE": 22,
-    "ENCOURAGING": 23,
-    "CONTEMPT": 24,
-}
-
-
-class TokenParser:
-    """Turn label to special token"""
-
-    def __init__(self):
-        pass
-
-    """Parse the attributes of a person."""
-
-    def __init__(self):
-        pass
-
-    @staticmethod
-    def age(age: str) -> str:
-        """Turn age token."""
-        age_id = AGE_MAP[age]
-        return f"<|age_{age_id}|>"
-
-    @staticmethod
-    def gender(gender: str) -> str:
-        """Turn gender token."""
-        gender_id = GENDER_MAP[gender]
-        return f"<|gender_{gender_id}|>"
-
-    @staticmethod
-    def mel_value(mel: int):
-        """Turn special token of mel scale pitch."""
-        mel = max(0, int(mel))
-        mel = min(1000, int(mel))
-        return f"<|pitch_value_{mel}|>"
-
-    @staticmethod
-    def mel_level(level: str):
-        """Turn special token of mel level."""
-        level_tag = LEVELS_MAP[level]
-        return f"<|pitch_label_{level_tag}|>"
-
-    @staticmethod
-    def pitch_var_value(pitch_std: int):
-        """Turn special token of pitch_std value."""
-        assert isinstance(pitch_std, int)
-        pitch_std = max(0, int(pitch_std))
-        pitch_std = min(10, int(pitch_std))
-        return f"<|pitch_var_value_{pitch_std}|>"
-
-    @staticmethod
-    def pitch_var_level(level: str):
-        """Turn special token of pitch std level."""
-        level_tag = LEVELS_MAP[level]
-        return f"<|pitch_var_label_{level_tag}|>"
-
-    @staticmethod
-    def loudness_value(loudness: int):
-        """Turn special toak of loudness value [0, 30]"""
-        assert loudness >= 0
-        loudness = max(0, int(loudness))
-        loudness = min(30, int(loudness))
-        return f"<|loudness_value_{loudness}|>"
-
-    @staticmethod
-    def loudness_level(level: str):
-        """Turn special token of loudness level."""
-        level_tag = LEVELS_MAP[level]
-        return f"<|loudness_label_{level_tag}|>"
-
-    @staticmethod
-    def speed_value(speed: int):
-        """Turn special token of speed value."""
-        speed = max(0, int(speed))
-        speed = min(10, int(speed))
-        return f"<|speed_value_{speed}|>"
-
-    @staticmethod
-    def speed_level(level: str):
-        """Turn special token of speed level."""
-        level_tag = LEVELS_MAP[level]
-        return f"<|speed_label_{level_tag}|>"
-
-    @staticmethod
-    def task(task: str) -> str:
-        """Turn special token of task."""
-        assert task in TASK_TOKEN_MAP.keys()
-
-        return TASK_TOKEN_MAP[task]
-
-    @staticmethod
-    def emotion(emotion: str):
-        emo_id = EMO_MAP[emotion]
-
-        return f"<|emotion_{emo_id}|>"
-
-# =============================================================================
-# >> END: PASTE CODE FROM sparktts/utils/* HERE <<
-# =============================================================================
-
-
-class SparkTTSProcessor(ProcessorMixin, PushToHubMixin): # Added PushToHubMixin
-    r"""
-    Constructs a SparkTTS processor which wraps a text tokenizer and relevant audio processing logic.
-
-    Args:
-        tokenizer ([`PreTrainedTokenizer`]):
-            An instance of [`PreTrainedTokenizer`]. This handles the text tokenization for the LLM.
-        feature_extractor ([`Wav2Vec2FeatureExtractor`]):
-            An instance of [`Wav2Vec2FeatureExtractor`]. Although Wav2Vec2 features are extracted
-            within the model's `tokenize_audio`, the extractor's configuration (like sampling rate)
-            is useful, and it aligns with the ProcessorMixin pattern.
-        config ([`SparkTTSConfig`], *optional*):
-            An instance of [`SparkTTSConfig`] to access configuration parameters like sample rate.
-    """
-    attributes = ["tokenizer", "feature_extractor"]
-    tokenizer_class = "AutoTokenizer"
-    feature_extractor_class = "Wav2Vec2FeatureExtractor" # Keep for consistency
-
-    def __init__(self, tokenizer, feature_extractor, config: Optional[SparkTTSConfig] = None, **kwargs):
-        super().__init__(tokenizer=tokenizer, feature_extractor=feature_extractor, **kwargs)
-        self.model = None
-        self.config = config
-        # Set sampling rate
-        if config and hasattr(config, 'sample_rate'):
-             self.sampling_rate = config.sample_rate
-        elif feature_extractor and hasattr(feature_extractor, 'sampling_rate'):
-             self.sampling_rate = feature_extractor.sampling_rate
-        else:
-             self.sampling_rate = 16000
-             logger.warning(f"Could not determine sampling rate. Defaulting to {self.sampling_rate} Hz.")
-
-        # # Ensure tokenizer pad token
-        # if self.tokenizer.pad_token is None:
-        #     if self.tokenizer.eos_token is not None:
-        #         logger.warning("Tokenizer does not have a pad token. Setting pad_token to eos_token.")
-        #         self.tokenizer.pad_token = self.tokenizer.eos_token
-        #     else:
-        #          logger.warning("Tokenizer lacks pad and eos token. Adding default pad token '<|pad|>'.")
-        #          self.tokenizer.add_special_tokens({'pad_token': '<|pad|>'})
-
-    def link_model(self, model):
-        """Links the processor to a SparkTTSModel instance for audio processing calls."""
-        if not hasattr(model, 'tokenize_audio') or not hasattr(model, 'detokenize_audio'):
-             raise TypeError("The provided model instance does not have the required 'tokenize_audio' and 'detokenize_audio' methods.")
-        if not hasattr(model, 'config'):
-             logger.warning("Linked model does not have a 'config' attribute. Some processor functionalities might rely on it.")
-
-        self.model = model
-        logger.info("SparkTTSModel successfully linked to the processor.")
-        # Update sampling rate based on linked model's config if available
-        if hasattr(model, 'config') and hasattr(model.config, 'sample_rate'):
-             if self.sampling_rate != model.config.sample_rate:
-                  logger.info(f"Updating processor sampling rate from {self.sampling_rate} to {model.config.sample_rate} based on linked model config.")
-                  self.sampling_rate = model.config.sample_rate
-             # Also update feature extractor sampling rate if it differs
-             if hasattr(self, 'feature_extractor') and self.feature_extractor.sampling_rate != model.config.sample_rate:
-                  logger.info(f"Updating feature_extractor sampling rate from {self.feature_extractor.sampling_rate} to {model.config.sample_rate}.")
-                  self.feature_extractor.sampling_rate = model.config.sample_rate
-
-
-    def __call__(
-        self,
-        text: str,
-        prompt_speech_path: Optional[Union[str, Path]] = None,
-        prompt_text: Optional[str] = None,
-        gender: Optional[str] = None,
-        pitch: Optional[str] = None,
-        speed: Optional[str] = None,
-        return_tensors: Optional[str] = "pt",
-        **kwargs, # Allow passing other args like padding, truncation to tokenizer
-    ) -> BatchEncoding:
-        """
-        Processes the input text and optional prompt audio/control parameters into a format suitable for [`SparkTTSModel`].
-
-        Args:
-            text (`str`):
-                The main text to be synthesized.
-            prompt_speech_path (`str` or `Path`, *optional*):
-                Path to the prompt audio file for voice cloning. Required if `gender` is not set.
-            prompt_text (`str`, *optional*):
-                Transcript of the prompt audio. Used only in voice cloning mode.
-            gender (`str`, *optional*):
-                Target gender ("male" or "female") for controllable synthesis. If set, enables control mode.
-            pitch (`str`, *optional*):
-                Target pitch level ("very_low", "low", "moderate", "high", "very_high") for control mode. Required if `gender` is set.
-            speed (`str`, *optional*):
-                Target speed level ("very_low", "low", "moderate", "high", "very_high") for control mode. Required if `gender` is set.
-            return_tensors (`str`, *optional*, defaults to `"pt"`):
-                If set, will return tensors instead of list of python integers. Only "pt" (PyTorch) is supported currently.
-            **kwargs:
-                Additional arguments passed to the underlying tokenizer's `__call__` method.
-
-        Returns:
-            [`BatchEncoding`]: A dictionary containing the `input_ids` and `attention_mask` for the LLM.
-            In voice cloning mode, it also includes `global_token_ids_prompt` (torch.Tensor) representing the
-            global tokens extracted from the prompt audio.
-        """
-
-        global_token_ids_prompt = None # Initialize
-
-        # Determine mode: Control TTS or Voice Cloning (Prompt TTS)
-        is_control_mode = gender is not None
-        is_cloning_mode = prompt_speech_path is not None and not is_control_mode
-
-        if is_control_mode:
-            # --- Controllable TTS Prompt Construction ---
-            if not all([pitch, speed]):
-                raise ValueError("For controllable TTS, 'gender', 'pitch', and 'speed' must all be provided.")
-            if prompt_speech_path is not None:
-                 logger.warning("`prompt_speech_path` provided but ignored because `gender` is set (controllable TTS mode).")
-
-            if not all(k in GENDER_MAP for k in [gender]): # Basic check
-                 raise ValueError(f"Invalid gender provided: {gender}. Must be one of {list(GENDER_MAP.keys())}")
-            if not all(k in LEVELS_MAP for k in [pitch, speed]): # Basic check
-                 raise ValueError(f"Invalid pitch or speed level provided. Must be one of {list(LEVELS_MAP.keys())}")
-
-            gender_id = GENDER_MAP[gender]
-            pitch_level_id = LEVELS_MAP[pitch]
-            speed_level_id = LEVELS_MAP[speed]
-
-            pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>"
-            speed_label_tokens = f"<|speed_label_{speed_level_id}|>"
-            gender_tokens = f"<|gender_{gender_id}|>"
-
-            attribute_tokens = "".join([gender_tokens, pitch_label_tokens, speed_label_tokens])
-
-            prompt_list = [
-                TASK_TOKEN_MAP["controllable_tts"],
-                "<|start_content|>",
-                text,
-                "<|end_content|>",
-                "<|start_style_label|>",
-                attribute_tokens,
-                "<|end_style_label|>",
-            ]
-            prompt_string = "".join(prompt_list)
-
-        elif is_cloning_mode:
-            # --- Voice Cloning Prompt Construction ---
-            if self.model is None:
-                raise RuntimeError("Processor must be linked to a SparkTTSModel instance via `processor.link_model(model)` before performing voice cloning.")
-            prompt_speech_path = Path(prompt_speech_path) # Ensure it's a Path object
-            if not prompt_speech_path.exists():
-                 raise FileNotFoundError(f"Prompt audio file not found: {prompt_speech_path}")
-
-            # Load and process prompt audio
-            try:
-                model_config = self.model.config if self.model and hasattr(self.model, 'config') else self.config
-                if model_config is None:
-                     raise ValueError("Configuration not available in processor or linked model.")
-
-                # Load main wav
-                wav = load_audio(
-                    prompt_speech_path,
-                    sampling_rate=self.sampling_rate,
-                    volume_normalize=getattr(model_config, 'volume_normalize', True), # Use getattr for safety
-                )
-                # Get reference clip
-                wav_ref_np = get_ref_clip(wav, model_config) # Pass config object
-                wav_ref = torch.from_numpy(wav_ref_np).unsqueeze(0).float()
-                wav_tensor = torch.from_numpy(wav).unsqueeze(0).float()
-
-                # Tokenize using the linked model's method
-                # Assuming tokenize_audio returns tensors with batch dim 1: [1, N_global], [1, N_semantic]
-                global_tokens_tensor, semantic_tokens_tensor = self.model.tokenize_audio(wav_tensor, wav_ref)
-
-                # Store the global tokens tensor (with batch dim) for the output dict
-                global_token_ids_prompt = global_tokens_tensor # Keep batch dim [1, N_global]
-
-                # Convert tensors to lists of ints for string formatting
-                global_token_list = global_tokens_tensor.squeeze().tolist() # Remove batch dim -> list
-                semantic_token_list = semantic_tokens_tensor.squeeze().tolist() # Remove batch dim -> list
-
-            except Exception as e:
-                logger.error(f"Error processing prompt audio {prompt_speech_path}: {e}")
-                import traceback
-                traceback.print_exc()
-                raise
-
-            # ==============================================================
-            # CORRECTED TOKEN STRING FORMATTING
-            # ==============================================================
-            # Create individual token strings for each ID
-            global_tokens_str = "".join([f"<|bicodec_global_{gid}|>" for gid in global_token_list])
-            semantic_tokens_str = "".join([f"<|bicodec_semantic_{sid}|>" for sid in semantic_token_list])
-            # ==============================================================
-            
-            # Construct prompt list based on presence of prompt_text
-            if prompt_text is not None and prompt_text.strip(): # Check if prompt_text is meaningful
-                logger.info("Using prompt text in voice cloning prompt.")
-                prompt_list = [
-                    TASK_TOKEN_MAP["tts"], # Or maybe TASK_TOKEN_MAP["prompt_tts"]? Check original logic. Assuming "tts".
-                    "<|start_content|>",
-                    prompt_text, # Transcript first
-                    text,        # Then target text
-                    "<|end_content|>",
-                    "<|start_global_token|>",
-                    global_tokens_str,
-                    "<|end_global_token|>",
-                    "<|start_semantic_token|>",
-                    semantic_tokens_str,
-                    # "<|end_semantic_token|>", # Original code didn't have this marker here
-                ]
-            else:
-                 # Simpler prompt without semantic tokens if no transcript provided
-                 logger.info("No prompt text provided, using text-only voice cloning prompt.")
-                 prompt_list = [
-                    TASK_TOKEN_MAP["tts"], # Or maybe TASK_TOKEN_MAP["prompt_tts"]?
-                    "<|start_content|>",
-                    text, # Only target text
-                    "<|end_content|>",
-                    "<|start_global_token|>",
-                    global_tokens_str,
-                    "<|end_global_token|>",
-                 ]
-            prompt_string = "".join(prompt_list)
-            logger.debug(f"Generated prompt string (cloning): {prompt_string[:200]}...") # Log start of prompt
-
-        else:
-            raise ValueError("Invalid input combination. Either provide `prompt_speech_path` for cloning or (`gender`, `pitch`, `speed`) for control.")
-
-        # --- Tokenize the final prompt string ---
-        # print(f"Tokenizing prompt: {prompt_string}")
-        inputs = self.tokenizer(
-            prompt_string,
-            return_tensors=return_tensors,
-            padding=kwargs.get("padding", False), # Often False for generation prompts unless batching > 1
-            truncation=kwargs.get("truncation", True),
-            max_length=kwargs.get("max_length", self.tokenizer.model_max_length),
-            add_special_tokens=kwargs.get("add_special_tokens", True), # Usually True unless handled manually
-            return_attention_mask=kwargs.get("return_attention_mask", True), # Need attention mask
-            **{k: v for k, v in kwargs.items() if k not in ["padding", "truncation", "max_length", "add_special_tokens", "return_attention_mask"]}
-        )
-        logger.debug(f"Tokenized input_ids shape: {inputs['input_ids'].shape}")
-
-
-        # Add the prompt's global tokens (as tensor with batch dim) to the output if in cloning mode
-        if is_cloning_mode and global_token_ids_prompt is not None:
-            if return_tensors == "pt":
-                 inputs["global_token_ids_prompt"] = global_token_ids_prompt # Already has batch dim [1, N_global]
-            else:
-                 # Handle non-tensor return if necessary
-                 inputs["global_token_ids_prompt"] = global_token_ids_prompt.tolist()
-
-        return inputs
-
-
-    def decode(
-        self,
-        generated_ids: torch.Tensor,
-        global_token_ids_prompt: Optional[torch.Tensor] = None,
-        input_ids_len: Optional[int] = None,
-        skip_special_tokens: bool = True,
-    ) -> Dict[str, Any]:
-        """
-        Decodes the generated token IDs from [`SparkTTSModel`] into an audio waveform.
-
-        Args:
-            generated_ids (`torch.Tensor`):
-                Tensor of token IDs generated by `model.generate()`, including the input prompt part. Shape [B, seq_len].
-            global_token_ids_prompt (`torch.Tensor`, *optional*):
-                The global tokens extracted from the prompt audio during the `__call__` step (for voice cloning).
-                Shape [B, N_global]. Required if the generation was for voice cloning.
-            input_ids_len (`int`, *optional*):
-                The length of the original input prompt `input_ids` fed to `model.generate()`. Required to
-                correctly isolate the newly generated tokens.
-            skip_special_tokens (`bool`, *optional*, defaults to `True`):
-                Whether to skip special tokens during the text decoding step (used to extract audio tokens).
-
-        Returns:
-            Dict[str, Any]: A dictionary containing:
-                - "audio": The decoded audio waveform as a NumPy array. Shape [T_audio] (if B=1) or [B, T_audio].
-                - "sampling_rate": The sampling rate of the audio.
-        """
-        if self.model is None:
-            raise RuntimeError("Processor must be linked to a SparkTTSModel instance via `processor.link_model(model)` before decoding.")
-        if input_ids_len is None:
-            raise ValueError("`input_ids_len` (length of the prompt input_ids) must be provided for decoding.")
-
-        # --- Isolate generated part and decode text ---
-        # Assumes generated_ids has shape [B, full_seq_len]
-        # Handle case where generated sequence is shorter than prompt (shouldn't happen with max_new_tokens > 0)
-        if generated_ids.shape[1] < input_ids_len:
-             logger.warning(f"Generated sequence length ({generated_ids.shape[1]}) is shorter than input prompt length ({input_ids_len}). Decoding might be incorrect.")
-             output_only_ids = generated_ids[:, input_ids_len:] # Will be empty if equal
-        else:
-             output_only_ids = generated_ids[:, input_ids_len:]
-
-
-        # Decode the generated part to find audio tokens
-        # Need to handle batch decoding if B > 1
-        # print("decode token", self.tokenizer.batch_decode(output_only_ids, skip_special_tokens=False))
-        decoded_texts = self.tokenizer.batch_decode(output_only_ids, skip_special_tokens=skip_special_tokens)
-
-        # --- Extract Audio Tokens ---
-        # Handle batch processing correctly
-        batch_size = generated_ids.shape[0]
-        all_semantic_ids = []
-        all_global_tokens = []
-        successful_indices = [] # Keep track of which batch items were successful
-
-        for i in range(batch_size):
-            decoded_text = decoded_texts[i]
-            current_semantic_ids = None
-            current_global_tokens = None
-
-            # Extract semantic tokens
-            try:
-                pred_semantic_indices = [int(token) for token in re.findall(r"bicodec_semantic_(\d+)", decoded_text)]
-                if not pred_semantic_indices:
-                    logger.warning(f"Batch item {i}: No semantic tokens found in decoded text: '{decoded_text[:200]}...'")
-                    continue # Skip this item
-
-                current_semantic_ids = torch.tensor(pred_semantic_indices).long() # Shape [N_semantic]
-            except Exception as e:
-                logger.error(f"Batch item {i}: Error parsing semantic tokens from: '{decoded_text[:200]}...'. Error: {e}")
-                continue # Skip this item
-
-            # Determine global tokens
-            if global_token_ids_prompt is not None:
-                # Cloning mode: Use the provided prompt global tokens for this batch item
-                if global_token_ids_prompt.shape[0] != batch_size:
-                     raise ValueError(f"Batch size mismatch: generated_ids has {batch_size}, but global_token_ids_prompt has {global_token_ids_prompt.shape[0]}.")
-                current_global_tokens = global_token_ids_prompt[i] # Shape [N_global]
-            else:
-                # Control mode: Extract global tokens from the generated text
-                try:
-                    pred_global_indices = [int(token) for token in re.findall(r"bicodec_global_(\d+)", decoded_text)]
-                    if not pred_global_indices:
-                        logger.warning(f"Batch item {i}: No global tokens found in decoded text for control mode: '{decoded_text[:200]}...'")
-                        continue # Skip this item
-
-                    current_global_tokens = torch.tensor(pred_global_indices).long() # Shape [N_global]
-
-                except Exception as e:
-                    logger.error(f"Batch item {i}: Error parsing global tokens from: '{decoded_text[:200]}...'. Error: {e}")
-                    continue # Skip this item
-
-            # If both tokens extracted successfully
-            all_semantic_ids.append(current_semantic_ids)
-            all_global_tokens.append(current_global_tokens)
-            successful_indices.append(i)
-
-        if not successful_indices:
-            logger.error("Failed to extract audio tokens for any item in the batch.")
-            return {"audio": np.array([], dtype=np.float32), "sampling_rate": self.sampling_rate}
-
-        # Pad sequences to the max length within the successful batch items for batch detokenization
-        # Note: BiCodec might not support batching if sequences have different lengths. Check its implementation.
-        # Assuming BiCodec *can* handle batches if padded (or if lengths are naturally equal).
-        # This padding might be unnecessary if BiCodec handles variable lengths or if B=1 anyway.
-        # For now, let's assume B=1 was handled correctly and skip complex padding.
-        if batch_size > 1 and len(successful_indices) < batch_size:
-             logger.warning(f"Only successfully decoded {len(successful_indices)} out of {batch_size} batch items.")
-             # Further processing might need to handle only the successful items.
-
-        # Let's proceed assuming B=1 or BiCodec handles batches appropriately.
-        # Stack the successful tokens.
-        try:
-            # Need to ensure tensors have the same length before stacking if BiCodec requires it.
-            # If BiCodec handles variable length, stacking might not be needed, just loop and call detokenize.
-            # Let's assume B=1 for simplicity of the example, matching original code's likely behavior.
-            if len(successful_indices) != 1:
-                 raise NotImplementedError("Batch decoding (B > 1) requires verification of BiCodec's batch handling and potentially padding.")
-
-            final_semantic_ids = all_semantic_ids[0].unsqueeze(0) # Add batch dim [1, N_semantic]
-            final_global_tokens = all_global_tokens[0].unsqueeze(0) # Add batch dim [1, N_global]
-
-        except IndexError: # Should not happen if successful_indices is not empty
-             logger.error("Internal error during token batch preparation.")
-             return {"audio": np.array([], dtype=np.float32), "sampling_rate": self.sampling_rate}
-
-
-        # --- Detokenize Audio ---
-        try:
-            # Call the linked model's detokenize method
-            # print(f"DEBUG: Detokenizing audio with global tokens {final_global_tokens.shape}, semantic tokens {final_semantic_ids.shape}")
-            output_wav = self.model.detokenize_audio(final_global_tokens, final_semantic_ids)
-            # detokenize_audio now returns numpy array float32 in [-1, 1]
-
-            # Optional: Double-check dtype here if needed, but should be handled by detokenize_audio now
-            # if output_wav.dtype != np.float32:
-            #    logger.warning(f"Audio dtype after detokenize is {output_wav.dtype}. Converting to float32.")
-            #    output_wav = output_wav.astype(np.float32)
-            # output_wav = np.clip(output_wav, -1.0, 1.0) # Clipping done in detokenize_audio
-
-        except Exception as e:
-             logger.error(f"Error during audio detokenization: {e}")
-             import traceback
-             traceback.print_exc()
-             raise RuntimeError("Audio detokenization failed.") from e
-
-        return {"audio": output_wav, "sampling_rate": self.sampling_rate}
-
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Union[str, os.PathLike],
-        cache_dir: Optional[Union[str, os.PathLike]] = None,
-        force_download: bool = False,
-        local_files_only: bool = False,
-        token: Optional[Union[str, bool]] = None,
-        revision: str = "main",
-        trust_remote_code: bool = False, # Allow passing this, needed for config potentially
-        **kwargs,
-    ):
-        r"""
-        Instantiate a SparkTTSProcessor from pretrained components.
-        """
-        # Pop specific kwargs for this method
-        config = kwargs.pop("config", None) # Allow passing config explicitly
-
-        # --- 1. Load Config (to find component paths) ---
-        # We need the config even if the processor doesn't store it permanently,
-        # just to find where the tokenizer/feature_extractor live.
-        loaded_config = None
-        if not isinstance(config, SparkTTSConfig):
-            try:
-                # Load the specific config class
-                loaded_config = SparkTTSConfig.from_pretrained(
-                    pretrained_model_name_or_path,
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    local_files_only=local_files_only,
-                    token=token,
-                    revision=revision,
-                    trust_remote_code=trust_remote_code, # Config might be custom
-                    **kwargs, # Pass relevant kwargs
-                )
-            except Exception as e:
-                logger.warning(
-                    f"Could not load SparkTTSConfig from {pretrained_model_name_or_path}. "
-                    f"Attempting to load components from default relative paths ('LLM', 'wav2vec2-large-xlsr-53'). Error: {e}"
-                )
-                loaded_config = None # Fallback
-        else:
-             # Config object was passed directly
-             loaded_config = config
-
-
-        # --- 2. Determine Component Paths ---
-        llm_tokenizer_path_or_id = "./LLM" # Default relative path
-        w2v_processor_path_or_id = "./wav2vec2-large-xlsr-53" # Default relative path
-
-        if loaded_config:
-            llm_tokenizer_path_or_id = getattr(loaded_config, 'llm_model_name_or_path', llm_tokenizer_path_or_id)
-            w2v_processor_path_or_id = getattr(loaded_config, 'wav2vec2_model_name_or_path', w2v_processor_path_or_id)
-
-        # The component `from_pretrained` methods handle resolving these paths/IDs
-        # whether they are relative subfolders of `pretrained_model_name_or_path`
-        # or separate Hub IDs.
-
-        # --- 3. Load Components ---
-        # Pass down relevant kwargs for loading components
-        component_loading_kwargs = {
-             "cache_dir": cache_dir,
-             "force_download": force_download,
-             "local_files_only": local_files_only,
-             "token": token,
-             "revision": revision,
-             **kwargs # Pass other user kwargs
-        }
-        try:
-            # Tokenizer might require trust_remote_code if its class is custom
-            tokenizer = AutoTokenizer.from_pretrained(
-                pretrained_model_name_or_path, # Main path
-                subfolder=llm_tokenizer_path_or_id.lstrip('./'), # Specify subfolder relative to main path
-                trust_remote_code=trust_remote_code,
-                **component_loading_kwargs
-            )
-        except Exception as e:
-            # Fallback: try loading directly using the path/id from config if different
-            if llm_tokenizer_path_or_id != "./LLM":
-                 try:
-                      logger.info(f"Retrying tokenizer load directly from: {llm_tokenizer_path_or_id}")
-                      tokenizer = AutoTokenizer.from_pretrained(
-                           llm_tokenizer_path_or_id,
-                           trust_remote_code=trust_remote_code,
-                           **component_loading_kwargs
-                      )
-                 except Exception as e2:
-                      raise OSError(f"Could not load tokenizer using main path + subfolder or directly from '{llm_tokenizer_path_or_id}'. Error: {e2}") from e
-            else:
-                 raise OSError(f"Could not load tokenizer from subfolder '{llm_tokenizer_path_or_id}' within '{pretrained_model_name_or_path}'. Error: {e}")
-
-
-        try:
-            # Feature extractor usually doesn't need trust_remote_code
-            feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
-                pretrained_model_name_or_path, # Main path
-                subfolder=w2v_processor_path_or_id.lstrip('./'), # Specify subfolder relative to main path
-                **component_loading_kwargs
-            )
-        except Exception as e:
-             # Fallback: try loading directly using the path/id from config if different
-            if w2v_processor_path_or_id != "./wav2vec2-large-xlsr-53":
-                 try:
-                      logger.info(f"Retrying feature extractor load directly from: {w2v_processor_path_or_id}")
-                      feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
-                           w2v_processor_path_or_id,
-                           **component_loading_kwargs
-                      )
-                 except Exception as e2:
-                      raise OSError(f"Could not load feature extractor using main path + subfolder or directly from '{w2v_processor_path_or_id}'. Error: {e2}") from e
-            else:
-                 raise OSError(f"Could not load feature extractor from subfolder '{w2v_processor_path_or_id}' within '{pretrained_model_name_or_path}'. Error: {e}")
-
-
-        # --- 4. Instantiate processor ---
-        # Pass the potentially loaded config object (or None)
-        return cls(tokenizer=tokenizer, feature_extractor=feature_extractor, config=loaded_config)
-
-    
-    def save_pretrained(
-        self,
-        save_directory: Union[str, os.PathLike],
-        push_to_hub: bool = False,
-        **kwargs,
-    ):
-        """
-        Save the processor's state (tokenizer and feature extractor files) to a directory.
-
-        Args:
-            save_directory (`str` or `os.PathLike`):
-                Directory where the processor files will be saved.
-            push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face Hub after saving it.
-            **kwargs:
-                Additional key word arguments passed along to the `push_to_hub` method.
-        """
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
-
-        # Save tokenizer
-        self.tokenizer.save_pretrained(str(save_directory), **kwargs)
-
-        # Save feature extractor
-        self.feature_extractor.save_pretrained(str(save_directory), **kwargs)
-
-        # Save the main processor config (if it exists and has relevant info)
-        # Note: The SparkTTSConfig is usually saved with the *model*, not the processor.
-        # However, if the processor holds specific config needed for reloading *itself*,
-        # it could be saved here. Usually, relying on the model's config is sufficient.
-        # if self.config:
-        #    self.config.save_pretrained(str(save_directory)) # Example if needed
-
-        logger.info(f"Processor components saved in {save_directory}")
-
-        if push_to_hub:
-             # Commit message and other hub kwargs can be passed via **kwargs
-             commit_message = kwargs.pop("commit_message", "Save processor")
-             return self.push_to_hub(save_directory, commit_message=commit_message, **kwargs)
-
-        return str(save_directory) # Return path consistent with Mixin