Update processing_gemma3_omni.py

processing_gemma3_omni.py

@@ -1,9 +1,9 @@
 import re
-from typing import List, Optional, Union, Dict, Any, Tuple
+from typing import List, Optional, Union, Dict, Any
 
 import math
 import numpy as np
-import scipy.signal  # Keep for resample_poly
+import scipy.signal
 import torch
 from torch.nn.utils.rnn import pad_sequence
 
@@ -15,28 +15,31 @@ AudioInput = Union[np.ndarray, List[float], Tuple[np.ndarray, int]]
 
 from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
 from transformers.feature_extraction_utils import BatchFeature
-# from transformers.image_utils import make_nested_list_of_images # Not used in this file
-from transformers.processing_utils import ProcessorMixin, \
-    ProcessingKwargs  # Removed ImagesKwargs, Unpack for simplicity as they are not fully defined here
+from transformers.processing_utils import ProcessorMixin, ProcessingKwargs 
 from transformers.utils import TensorType, to_py_obj, logging
+# For AutoImageProcessor, AutoTokenizer if needed for default loading
+from transformers import AutoImageProcessor, AutoTokenizer
 
-# Constants
+
+# Constants (as defined before)
 DEFAULT_SAMPLING_RATE = 16000
 DEFAULT_N_FFT = 512
-DEFAULT_WIN_LENGTH = 400
-DEFAULT_HOP_LENGTH = 160
+DEFAULT_WIN_LENGTH = 400 # Will be n_fft if None in __init__
+DEFAULT_HOP_LENGTH = 160 # Will be win_length // 4 if None in __init__
 DEFAULT_N_MELS = 80
-DEFAULT_COMPRESSION_RATE = 4  # As in FeatureExtractor init
-DEFAULT_QFORMER_RATE = 2  # As in FeatureExtractor init
-DEFAULT_FEAT_STRIDE = 4  # As in FeatureExtractor init, affects 'frames' if used
-IMAGE_TOKEN_PATTERN = r"<\|image_\d+\|>"  # From user
-AUDIO_TOKEN_PATTERN = r"<\|audio_\d+\|>"  # From user
-DEFAULT_MAX_LENGTH = 16384  # From user
+DEFAULT_COMPRESSION_RATE = 4
+DEFAULT_QFORMER_RATE = 2
+DEFAULT_FEAT_STRIDE = 4
+IMAGE_TOKEN_PATTERN = r"<\|image_\d+\|>"
+AUDIO_TOKEN_PATTERN = r"<\|audio_\d+\|>"
+DEFAULT_MAX_LENGTH = 16384
 LOG_MEL_CLIP_EPSILON = 1e-5
 
 logger = logging.get_logger(__name__)
 
 
+# create_mel_filterbank function (assuming it's correctly defined from previous response)
+# ... (create_mel_filterbank function from the previous corrected response) ...
 def create_mel_filterbank(sampling_rate: int, n_fft: int, n_mels: int, fmin: float = 0.0,
                           fmax: Optional[float] = None) -> np.ndarray:
     """Create Mel filterbank for audio processing."""
@@ -45,79 +48,66 @@ def create_mel_filterbank(sampling_rate: int, n_fft: int, n_mels: int, fmin: flo
     if fmin >= fmax:
         raise ValueError(f"fmin ({fmin}) must be smaller than fmax ({fmax}).")
 
-    def hz_to_mel(f: float) -> float:  # Using HTK formula (as in librosa default)
+    def hz_to_mel(f: float) -> float: # Using HTK formula (as in librosa default)
         return 2595.0 * math.log10(1 + f / 700.0)
 
     def mel_to_hz(mel: float) -> float:
-        return 700.0 * (10 ** (mel / 2595.0) - 1)
+        return 700.0 * (10**(mel / 2595.0) - 1)
 
     mel_points = np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2)
     freq_points = mel_to_hz(mel_points)
-
-    # Ensure freq_points are within the Nyquist limit
+    
     freq_points = np.clip(freq_points, 0, sampling_rate / 2.0)
-
-    # fftfreqs = np.fft.rfftfreq(n=n_fft, d=1.0 / sampling_rate) # Frequencies for each FFT bin center
-    # bins = np.searchsorted(fftfreqs, freq_points) # More robust way to find bins
-    bins = np.floor((n_fft / 2.0) * freq_points / (sampling_rate / 2.0)).astype(int)  # Simplified from librosa
-    bins = np.clip(bins, 0, n_fft // 2)  # Max index for rfft output (n_fft//2 + 1 bins, so max index is n_fft//2)
+    bins = np.floor((n_fft / 2.0) * freq_points / (sampling_rate / 2.0)).astype(int)
+    bins = np.clip(bins, 0, n_fft // 2)
 
     filterbank = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
-    for m in range(n_mels):  # Iterate 0 to n_mels-1 for m-th filter
+    for m in range(n_mels):
         left, center, right = bins[m], bins[m + 1], bins[m + 2]
-
+        
+        # Simplified triangle creation logic (more robust versions exist in libraries like librosa)
         if center > left:
-            filterbank[m, left:center + 1] = (np.arange(left, center + 1) - left) / (
-                        center - left)  # Inclusive center for peak
+            filterbank[m, left:center+1] = (np.arange(left, center + 1) - left) / (center - left)
         if right > center:
-            # The peak is at 'center'. So the downward slope starts from 'center'.
-            filterbank[m, center:right + 1] = (right - np.arange(center, right + 1)) / (
-                        right - center)  # Inclusive center
-        # Ensure the peak of the triangle is 1, and handle cases where left=center or center=right
-        # This simplified version might need librosa.filters.mel for full robustness if edge cases are hit.
-        # For now, ensuring distinct points for slopes:
-        if center > left and center < right:  # Standard triangle
-            # Ramp up
-            idxs_up = np.arange(left, center + 1)  # include center
-            filterbank[m, idxs_up] = (idxs_up - left) / (center - left)
-            # Ramp down
-            idxs_down = np.arange(center, right + 1)  # include center
-            filterbank[m, idxs_down] = (right - idxs_down) / (right - center)
-        elif center > left:  # only ramp up (right part is flat or non-existent)
-            idxs_up = np.arange(left, center + 1)
-            filterbank[m, idxs_up] = (idxs_up - left) / (center - left)
-        elif center < right:  # only ramp down (left part is flat or non-existent)
-            idxs_down = np.arange(center, right + 1)
-            filterbank[m, idxs_down] = (right - idxs_down) / (right - center)
+            filterbank[m, center:right+1] = (right - np.arange(center, right + 1)) / (right - center)
+        # Ensure peak is 1 if multiple points coincide at center (can happen with narrow filters/low resolution)
+        if left <= center <= right and filterbank[m,center] < 1.0 and (center > left or center < right) : #check if it's a valid point for a peak
+            # if filterbank[m,center] is not properly set to 1 by slopes (e.g. left==center or right==center)
+             filterbank[m,center] = 1.0
+             if left == center and right > center : # only falling slope
+                 # Ensure it doesn't double-dip if already set
+                 pass
+             elif right == center and left < center: # only rising slope
+                 pass
 
-    return filterbank
 
+    return filterbank
 
+# Gemma3AudioFeatureExtractor class (assuming it's correctly defined from previous response)
+# ... (Gemma3AudioFeatureExtractor class from the previous corrected response) ...
 class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
-    model_input_names = ["audio_values", "audio_attention_mask"]  # What this extractor produces for the model
+    model_input_names = ["audio_values", "audio_attention_mask"] 
 
     def __init__(
             self,
             compression_rate: int = DEFAULT_COMPRESSION_RATE,
             qformer_rate: int = DEFAULT_QFORMER_RATE,
             feat_stride: int = DEFAULT_FEAT_STRIDE,
-            sampling_rate: int = DEFAULT_SAMPLING_RATE,  # Target sampling rate
+            sampling_rate: int = DEFAULT_SAMPLING_RATE, 
             n_fft: int = DEFAULT_N_FFT,
-            win_length: Optional[int] = None,  # Default to n_fft if None
-            hop_length: Optional[int] = None,  # Default to win_length // 4 if None
+            win_length: Optional[int] = None, 
+            hop_length: Optional[int] = None, 
             n_mels: int = DEFAULT_N_MELS,
             f_min: float = 0.0,
             f_max: Optional[float] = None,
-            padding_value: float = 0.0,  # For pad_sequence of mels
+            padding_value: float = 0.0, 
             **kwargs
     ):
-        # feature_size is the number of features per frame (n_mels)
         super().__init__(feature_size=n_mels, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
 
         self.compression_rate = compression_rate
         self.qformer_rate = qformer_rate
-        self.feat_stride = feat_stride  # Used for 'frames' calculation, purpose might be downstream
-
+        self.feat_stride = feat_stride
         self.n_fft = n_fft
         self.win_length = win_length if win_length is not None else n_fft
         self.hop_length = hop_length if hop_length is not None else self.win_length // 4
@@ -130,26 +120,24 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
                 f"win_length ({self.win_length}) is greater than n_fft ({self.n_fft}). "
                 f"For FFT computation, the window will effectively be truncated or the signal zero-padded to n_fft length."
             )
-        self.window = scipy.signal.get_window("hann", self.win_length).astype(np.float32)  # Using Hann window
+        self.window = scipy.signal.get_window("hann", self.win_length).astype(np.float32)
         self.mel_filterbank = create_mel_filterbank(
             self.sampling_rate, self.n_fft, self.n_mels, fmin=self.f_min, fmax=self.f_max
-        ).T  # Transpose to (n_fft // 2 + 1, n_mels)
+        ).T
 
     def __call__(
             self,
             audios: Union[AudioInput, List[AudioInput]],
-            sampling_rate: Optional[int] = None,  # SR of the input audios if they are raw arrays
+            sampling_rate: Optional[int] = None, 
             return_tensors: Union[TensorType, str, None] = TensorType.PYTORCH
     ) -> BatchFeature:
-
+        
         if not isinstance(audios, list):
             audios = [audios]
 
         processed_mel_spectrograms: List[torch.Tensor] = []
         actual_mel_lengths: List[int] = []
-
-        # Optional downstream calculation values, kept if needed by other parts of Gemma3OmniProcessor
-        downstream_sizes_for_token_calc: List[torch.Tensor] = []
+        downstream_sizes_for_token_calc: List[torch.Tensor] = [] 
         downstream_frames_scaled_for_token_calc: List[int] = []
 
         for audio_input_item in audios:
@@ -166,73 +154,58 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
                         "sampling_rate must be provided if audio inputs are raw numpy arrays or lists."
                     )
                 source_sr = sampling_rate
-            else:  # Add more robust loading here if using file paths or bytes
+            else:
                 raise TypeError(
                     f"Unsupported audio input type: {type(audio_input_item)}. "
                     "This extractor expects np.ndarray, list of floats, or Tuple[np.ndarray, int indicating SR]."
                 )
 
             processed_wav = self._preprocess_audio(current_wav_array, source_sr)
-            # mel_spectrogram has shape (frame_count, n_mels)
-            mel_spectrogram = self._compute_log_mel_spectrogram(processed_wav)
-
-            feature_tensor = torch.from_numpy(mel_spectrogram)  # Already float32 from _compute_log_mel_spectrogram
+            mel_spectrogram = self._compute_log_mel_spectrogram(processed_wav) 
+            
+            feature_tensor = torch.from_numpy(mel_spectrogram)
             processed_mel_spectrograms.append(feature_tensor)
-            actual_mel_lengths.append(feature_tensor.shape[0])  # frame_count for this item
+            actual_mel_lengths.append(feature_tensor.shape[0]) 
 
-            # These calculations seem related to determining the number of special tokens in the prompt
             downstream_sizes_for_token_calc.append(torch.tensor(self._calculate_embed_length(feature_tensor.shape[0])))
             downstream_frames_scaled_for_token_calc.append(feature_tensor.shape[0] * self.feat_stride)
 
-        # Pad the mel spectrograms
-        # audio_values will have shape (Batch, Max_frame_count, n_mels)
         audio_values = pad_sequence(processed_mel_spectrograms, batch_first=True, padding_value=self.padding_value)
-
-        # Create attention mask corresponding to audio_values
-        max_mel_len = audio_values.shape[1]  # Max_frame_count across the batch
+        max_mel_len = audio_values.shape[1]
         lengths_tensor = torch.tensor(actual_mel_lengths, dtype=torch.long)
-
-        audio_attention_mask = torch.arange(max_mel_len).unsqueeze(0).expand(len(audios),
-                                                                             -1) < lengths_tensor.unsqueeze(1)
-
+        audio_attention_mask = torch.arange(max_mel_len).unsqueeze(0).expand(len(audios), -1) < lengths_tensor.unsqueeze(1)
+        
         output_data = {
-            "audio_values": audio_values,  # (B, Max_T_mel, N_Mels) -> Input to Conformer
-            "audio_attention_mask": audio_attention_mask  # (B, Max_T_mel) -> Mask for Conformer input
+            "audio_values": audio_values,
+            "audio_attention_mask": audio_attention_mask
         }
-
-        # Include these if they are used by Gemma3OmniProcessor for prompt construction
-        if downstream_sizes_for_token_calc:  # Renamed to clarify purpose
-            output_data["audio_token_calc_sizes"] = torch.stack(downstream_sizes_for_token_calc)
-        # if downstream_frames_scaled_for_token_calc: # Example if needed
-        #    output_data["audio_token_calc_frames_scaled"] = torch.tensor(downstream_frames_scaled_for_token_calc)
-
+        
+        if downstream_sizes_for_token_calc:
+             output_data["audio_token_calc_sizes"] = torch.stack(downstream_sizes_for_token_calc)
+        
         return BatchFeature(data=output_data, tensor_type=return_tensors)
 
     def _preprocess_audio(self, wav: np.ndarray, source_sr: int) -> np.ndarray:
         if wav.dtype not in [np.float32, np.float64]:
-            # Assuming int audio needs normalization to [-1, 1]
             if np.issubdtype(wav.dtype, np.integer):
                 max_val = np.iinfo(wav.dtype).max
                 wav = wav.astype(np.float32) / max_val
-            else:  # Other non-float types
+            else: 
                 wav = wav.astype(np.float32)
-
+        
         if wav.ndim > 1:
-            wav = wav.mean(axis=0)  # Convert to mono
-
+            wav = wav.mean(axis=0)
+        
         if source_sr != self.sampling_rate:
-            # Using scipy.signal.resample_poly for potentially higher quality resampling
-            # It requires integer up/down factors.
             gcd = math.gcd(self.sampling_rate, source_sr)
             up_factor = self.sampling_rate // gcd
             down_factor = source_sr // gcd
-            if up_factor != down_factor:  # Only resample if ratio is not 1
-                logger.info(f"Resampling audio from {source_sr} Hz to {self.sampling_rate} Hz.")
-                wav = scipy.signal.resample_poly(wav, up=up_factor, down=down_factor)
-
-        # Peak normalization
+            if up_factor != down_factor: 
+                 logger.info(f"Resampling audio from {source_sr} Hz to {self.sampling_rate} Hz.")
+                 wav = scipy.signal.resample_poly(wav, up=up_factor, down=down_factor)
+        
         norm_factor = np.abs(wav).max()
-        if norm_factor > 1e-9:  # Avoid division by zero/small numbers for silent/near-silent audio
+        if norm_factor > 1e-9: 
             wav = wav / norm_factor
         return wav
 
@@ -241,12 +214,9 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
             padding = self.win_length - len(wav)
             wav = np.pad(wav, (0, padding), mode='constant', constant_values=0.0)
 
-        # Using librosa-like STFT parameters where n_fft, hop_length, win_length are explicit
-        # Manual framing and windowing:
         num_frames = 1 + (len(wav) - self.win_length) // self.hop_length
         if num_frames <= 0:
-            logger.warning(
-                f"Audio of length {len(wav)} is too short to produce frames with win_length {self.win_length} and hop_length {self.hop_length}. Returning empty mel spectrogram.")
+            logger.warning(f"Audio of length {len(wav)} is too short to produce frames with win_length {self.win_length} and hop_length {self.hop_length}. Returning empty mel spectrogram.")
             return np.zeros((0, self.n_mels), dtype=np.float32)
 
         frames = np.lib.stride_tricks.as_strided(
@@ -255,153 +225,144 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
             strides=(wav.strides[0] * self.hop_length, wav.strides[0]),
             writeable=False
         )
-
+        
         windowed_frames = frames * self.window
-
-        stft_matrix = np.fft.rfft(windowed_frames, n=self.n_fft, axis=-1)  # Shape (num_frames, n_fft // 2 + 1)
-        powers = np.abs(stft_matrix) ** 2
-
-        mel_spectrogram = np.dot(powers, self.mel_filterbank)  # Shape (num_frames, n_mels)
-
+        stft_matrix = np.fft.rfft(windowed_frames, n=self.n_fft, axis=-1)
+        powers = np.abs(stft_matrix)**2
+        mel_spectrogram = np.dot(powers, self.mel_filterbank)
         mel_spectrogram = np.clip(mel_spectrogram, LOG_MEL_CLIP_EPSILON, None)
         log_mel_spectrogram = np.log(mel_spectrogram)
-
+        
         return log_mel_spectrogram.astype(np.float32)
 
     def _calculate_embed_length(self, frame_count: int) -> int:
-        # This calculation is likely for determining the number of special tokens
-        # to insert in the text prompt, based on the audio length.
         compressed = math.ceil(frame_count / self.compression_rate)
         return math.ceil(compressed / self.qformer_rate)
 
-
-# --- Gemma3ProcessorKwargs and Gemma3ImagesKwargs would be defined here if needed ---
-# For this fix, focusing on Gemma3AudioFeatureExtractor and Gemma3OmniProcessor interactions
-class Gemma3DummyProcessorKwargs(ProcessingKwargs, total=False):  # Dummy for testing
+class Gemma3DummyProcessorKwargs(ProcessingKwargs, total=False): # Dummy for testing structure
     images_kwargs: Dict[str, Any]
     audio_kwargs: Dict[str, Any]
-    text_kwargs: Dict[str, Any]  # Added for completeness
+    text_kwargs: Dict[str, Any]
     _defaults = {
         "text_kwargs": {"padding": False, "truncation": False, "max_length": DEFAULT_MAX_LENGTH},
         "images_kwargs": {},
         "audio_kwargs": {}
     }
 
-
 class Gemma3OmniProcessor(ProcessorMixin):
     attributes = ["image_processor", "audio_processor", "tokenizer"]
-    # Define image_processor_class and tokenizer_class if they are standard,
-    # or handle their instantiation/passing carefully.
-    # For custom audio_processor, we reference the class directly.
-    audio_processor_class = Gemma3AudioFeatureExtractor
+    # Define class attributes for ProcessorMixin to find/use them
+    image_processor_class = "AutoImageProcessor"  # Or the specific class string if not auto
+    audio_processor_class = Gemma3AudioFeatureExtractor # Correctly points to your custom class
+    tokenizer_class = "AutoTokenizer" # Or the specific class string
 
-    # image_processor_class = "AutoImageProcessor" # Example
-    # tokenizer_class = "AutoTokenizer" # Example
+    # valid_kwargs was in user's code, its role depends on ProcessorMixin internal usage
+    valid_kwargs = ["chat_template", "image_seq_length"] 
 
     def __init__(
             self,
-            tokenizer,  # Tokenizer is essential
+            tokenizer, 
             audio_processor: Optional[Union[Gemma3AudioFeatureExtractor, Dict]] = None,
-            image_processor=None,  # Define further if used
+            image_processor = None, 
             chat_template=None,
-            image_seq_length: int = 256,  # Default from user code
-            # Parameters for calculating number of audio soft tokens in text prompt
-            audio_prompt_compression_rate: int = 8,
+            image_seq_length: int = 256,
+            audio_prompt_compression_rate: int = 8, 
             audio_prompt_qformer_rate: int = 1,
             audio_prompt_feat_stride: int = 1,
-            audio_placeholder_token: str = "<|audio_placeholder|>",  # Placeholder in user text
-            audio_soft_token_str: str = "<audio_soft_token>",  # The actual soft token string
+            audio_placeholder_token: str = "<|audio_placeholder|>",
+            audio_soft_token_str: str = "<audio_soft_token>",
             **kwargs
     ):
+        # Instantiate audio_processor if config dict is passed or if None (use defaults)
         if audio_processor is None:
-            logger.info("Initializing Gemma3AudioFeatureExtractor with default parameters.")
+            logger.info("Initializing Gemma3AudioFeatureExtractor with default parameters for Gemma3OmniProcessor.")
             audio_processor = Gemma3AudioFeatureExtractor()
         elif isinstance(audio_processor, Dict):
             audio_processor = Gemma3AudioFeatureExtractor(**audio_processor)
-        elif not isinstance(audio_processor, Gemma3AudioFeatureExtractor):
-            raise TypeError(
-                f"audio_processor must be an instance of Gemma3AudioFeatureExtractor or a config dict, got {type(audio_processor)}")
+        elif not isinstance(audio_processor, Gemma3AudioFeatureExtractor): # Check type if instance is passed
+            raise TypeError(f"audio_processor must be an instance of Gemma3AudioFeatureExtractor or a config dict, got {type(audio_processor)}")
+
+        # Handle image_processor similarly if it can be None or a dict
+        if image_processor is None and self.image_processor_class:
+             # This is a basic way; from_pretrained usually handles complex loading
+            if isinstance(self.image_processor_class, str) and self.image_processor_class == "AutoImageProcessor":
+                logger.info(f"Attempting to load a default {self.image_processor_class}. This might require a default model name or fail.")
+                # image_processor = AutoImageProcessor.from_pretrained("openai/clip-vit-base-patch32") # Example default
+            # else if self.image_processor_class is an actual class, instantiate it.
+        elif isinstance(image_processor, Dict):
+            # image_processor = AutoImageProcessor.from_config(config_class(**image_processor)) # Example
+            pass # Actual instantiation from dict would be more complex
+
+        # Ensure tokenizer is an instantiated object
+        if isinstance(tokenizer, str): # If tokenizer is a string (model name/path)
+            logger.info(f"Loading tokenizer from {tokenizer}")
+            # tokenizer = AutoTokenizer.from_pretrained(tokenizer) # This is how it's usually done
+        elif tokenizer is None:
+             raise ValueError("A tokenizer instance or identifier must be provided.")
 
-        # Ensure tokenizer is provided and instantiated
-        if tokenizer is None:  # This check might be redundant if from_pretrained handles it
-            raise ValueError("A tokenizer must be provided.")
-        # if isinstance(tokenizer, str): # Basic loading, usually from_pretrained handles complex cases
-        #    tokenizer = AutoTokenizer.from_pretrained(tokenizer)
 
         super().__init__(
             image_processor=image_processor,
             audio_processor=audio_processor,
             tokenizer=tokenizer,
             chat_template=chat_template,
-            **kwargs
+            **kwargs # Pass other kwargs to super
         )
-
+        
         self.image_seq_length = image_seq_length
-        # Robustly get special image tokens from tokenizer
-        self.image_token_id = getattr(tokenizer, "image_token_id",
-                                      tokenizer.unk_token_id if hasattr(tokenizer, "unk_token_id") else None)
-        self.boi_token = getattr(tokenizer, "boi_token", "<|image|>")  # Using <|image|> as a more common BOI
-        self.image_token = getattr(tokenizer, "image_token", "<|image|>")
-        self.eoi_token = getattr(tokenizer, "eoi_token", "")  # End of image, can be empty
+        self.image_token_id = getattr(self.tokenizer, "image_token_id", self.tokenizer.unk_token_id if hasattr(self.tokenizer, "unk_token_id") else None)
+        self.boi_token = getattr(self.tokenizer, "boi_token", "<|image|>") 
+        self.image_token = getattr(self.tokenizer, "image_token", "<|image|>")
+        self.eoi_token = getattr(self.tokenizer, "eoi_token", "") 
 
         self.audio_placeholder_token = audio_placeholder_token
         self.audio_soft_token_str = audio_soft_token_str
-
-        # Get ID for the audio soft token string; it must exist in the tokenizer
+        
         self.audio_soft_token_id = self.tokenizer.convert_tokens_to_ids(self.audio_soft_token_str)
-        if self.audio_soft_token_id == self.tokenizer.unk_token_id:
-            logger.warning(
-                f"The audio soft token string '{self.audio_soft_token_str}' maps to UNK token. "
+        if self.audio_soft_token_id == self.tokenizer.unk_token_id: # Check if UNK
+             logger.warning(
+                f"The audio soft token string '{self.audio_soft_token_str}' maps to UNK token (ID: {self.audio_soft_token_id}). "
                 "Ensure it is added to the tokenizer's vocabulary as a special token."
             )
-        # User's original expected ID, for reference or potential validation
-        # self.expected_audio_token_id = 262143
-        # if self.audio_soft_token_id != self.expected_audio_token_id:
-        #     logger.warning(f"Assigned ID {self.audio_soft_token_id} for '{self.audio_soft_token_str}' does not match expected ID {self.expected_audio_token_id}.")
 
         self.full_image_sequence_str = f"\n\n{self.boi_token}{''.join([self.image_token] * self.image_seq_length)}{self.eoi_token}\n\n"
 
-        # Store rates for calculating number of audio soft tokens for the prompt
         self.audio_prompt_compression_rate = audio_prompt_compression_rate
         self.audio_prompt_qformer_rate = audio_prompt_qformer_rate
         self.audio_prompt_feat_stride = audio_prompt_feat_stride
 
-    def _merge_kwargs(self, ModelProcessorKwargs, tokenizer_init_kwargs, **kwargs_passed_to_call):
-        # Use ModelProcessorKwargs (e.g. Gemma3DummyProcessorKwargs) to get default structures
-        # This method was complex in user code, simplifying slightly for clarity
+
+    def _merge_kwargs(self, KwargsClassWithDefaults, tokenizer_init_kwargs, **kwargs_passed_to_call):
         final_kwargs = {}
         # Initialize with _defaults from the Kwargs class
-        for modality_key, default_modality_kwargs in ModelProcessorKwargs._defaults.items():
+        # Ensure KwargsClassWithDefaults has a _defaults attribute
+        _defaults = getattr(KwargsClassWithDefaults, "_defaults", {})
+        for modality_key, default_modality_kwargs in _defaults.items():
             final_kwargs[modality_key] = default_modality_kwargs.copy()
 
         # Override with tokenizer's init_kwargs if they exist for a given key
         for modality_key, modality_dict in final_kwargs.items():
-            for key in list(modality_dict.keys()):  # Iterate over copy of keys
+            for key in list(modality_dict.keys()): 
                 if key in tokenizer_init_kwargs:
                     modality_dict[key] = tokenizer_init_kwargs[key]
-
-        # Override with kwargs passed directly to __call__ (e.g. kwargs['text_kwargs'])
-        for modality_key, modality_dict_from_call in kwargs_passed_to_call.items():
-            if modality_key in final_kwargs and isinstance(modality_dict_from_call, dict):
-                final_kwargs[modality_key].update(modality_dict_from_call)
-            elif modality_key not in final_kwargs and isinstance(modality_dict_from_call,
-                                                                 dict):  # For kwargs not in _defaults
-                final_kwargs[modality_key] = modality_dict_from_call.copy()
-
-        # Specific handling for text_kwargs as in user code
-        if "text_kwargs" in final_kwargs:
-            final_kwargs["text_kwargs"]["truncation"] = final_kwargs["text_kwargs"].get("truncation",
-                                                                                        False)  # Keep user's if provided
-            final_kwargs["text_kwargs"]["max_length"] = final_kwargs["text_kwargs"].get("max_length",
-                                                                                        DEFAULT_MAX_LENGTH)
-        else:  # Ensure text_kwargs exists
-            final_kwargs["text_kwargs"] = {"padding": False, "truncation": False, "max_length": DEFAULT_MAX_LENGTH}
-
+        
+        # Override with kwargs passed directly to __call__
+        for modality_key_from_call, modality_dict_from_call in kwargs_passed_to_call.items():
+            if modality_key_from_call in final_kwargs and isinstance(modality_dict_from_call, dict):
+                final_kwargs[modality_key_from_call].update(modality_dict_from_call)
+            # If a new modality_kwargs (e.g., "video_kwargs") is passed, add it
+            elif modality_key_from_call not in final_kwargs and isinstance(modality_dict_from_call, dict):
+                 final_kwargs[modality_key_from_call] = modality_dict_from_call.copy()
+
+        # Specific handling for text_kwargs
+        if "text_kwargs" not in final_kwargs:
+             final_kwargs["text_kwargs"] = {} # Ensure it exists
+        final_kwargs["text_kwargs"]["truncation"] = final_kwargs["text_kwargs"].get("truncation", False)
+        final_kwargs["text_kwargs"]["max_length"] = final_kwargs["text_kwargs"].get("max_length", DEFAULT_MAX_LENGTH)
+            
         return final_kwargs
 
     def _compute_audio_prompt_token_count(self, actual_mel_frames_count: int) -> int:
-        """Calculates how many <audio_soft_token> to insert in the text prompt."""
-        # Uses parameters specific to this processor for prompt engineering
         scaled_frames = actual_mel_frames_count * self.audio_prompt_feat_stride
         compressed_once = math.ceil(scaled_frames / self.audio_prompt_compression_rate)
         compressed_twice = math.ceil(compressed_once / self.audio_prompt_qformer_rate)
@@ -410,118 +371,109 @@ class Gemma3OmniProcessor(ProcessorMixin):
     def __call__(
             self,
             text: Union[str, List[str]] = None,
-            images: Optional[Any] = None,  # Type depends on image_processor
+            images: Optional[Any] = None, 
             audios: Optional[Union[AudioInput, List[AudioInput]]] = None,
-            sampling_rate: Optional[int] = None,  # For raw audio arrays passed to audio_processor
-            return_tensors: Optional[Union[str, TensorType]] = None,  # Default behavior based on HF
-            **kwargs: Any  # Using Any for Unpack[Gemma3ProcessorKwargs] as it's not fully defined
+            sampling_rate: Optional[int] = None, 
+            return_tensors: Optional[Union[str, TensorType]] = None, 
+            **kwargs: Any 
     ) -> BatchFeature:
-
+        
         if text is None and images is None and audios is None:
             raise ValueError("Provide at least one of `text`, `images`, or `audios`.")
 
-        # Determine final return_tensors type (passed explicitly, or from text_kwargs, or default)
-        # The _merge_kwargs will get text_kwargs, but return_tensors can be a top-level arg.
-        _text_kwargs_from_call = kwargs.get("text_kwargs", {})
-        _rt_from_text_kwargs = _text_kwargs_from_call.get("return_tensors")
-
-        final_rt = TensorType.PYTORCH  # Default
-        if return_tensors is not None:
-            final_rt = return_tensors
-        elif _rt_from_text_kwargs is not None:
-            final_rt = _rt_from_text_kwargs
-
-        # Get all kwargs merged (text_kwargs, images_kwargs, audio_kwargs)
-        # Using Gemma3DummyProcessorKwargs as the source of _defaults structure
-        merged_kwargs = self._merge_kwargs(
-            Gemma3DummyProcessorKwargs,
+        # Determine final return_tensors strategy
+        # Priority: 1. Explicit return_tensors, 2. from text_kwargs in **kwargs, 3. Default (PT)
+        final_rt = return_tensors
+        merged_call_kwargs = self._merge_kwargs(
+            Gemma3DummyProcessorKwargs, # Using dummy for _defaults structure
             self.tokenizer.init_kwargs if hasattr(self.tokenizer, 'init_kwargs') else {},
-            **kwargs
+            **kwargs 
         )
+        
+        if final_rt is None: # If not passed directly to __call__
+            final_rt = merged_call_kwargs.get("text_kwargs", {}).pop("return_tensors", TensorType.PYTORCH)
+        else: # If passed directly, remove from text_kwargs to avoid conflict
+            merged_call_kwargs.get("text_kwargs", {}).pop("return_tensors", None)
+
 
-        # Ensure text is a list of strings
         if text is None:
             num_samples = 0
             if images is not None:
-                # Simplified: make_nested_list_of_images would be here for HF standard image proc.
-                num_samples = len(images) if isinstance(images, list) and not isinstance(images[0], (float, int)) else 1
+                _images_list = images if isinstance(images, list) and (not images or not isinstance(images[0], (int, float))) else [images]
+                num_samples = len(_images_list)
             elif audios is not None:
-                num_samples = len(audios) if isinstance(audios, list) else 1
-            text = [""] * num_samples if num_samples > 0 else [""]  # Dummy text if only modality
-
+                _audios_list = audios if isinstance(audios, list) else [audios]
+                num_samples = len(_audios_list)
+            text = [""] * num_samples if num_samples > 0 else [""]
+        
         if isinstance(text, str):
             text = [text]
         if not (isinstance(text, list) and all(isinstance(t, str) for t in text)):
             raise ValueError("Input `text` must be a string or a list of strings.")
 
-        # --- Image Processing (Simplified Placeholder) ---
-        image_features = {}
+        image_features_dict = {}
         if images is not None and self.image_processor is not None:
             logger.info("Processing images...")
-            # Actual image processing call would be here, e.g.:
-            # image_features = self.image_processor(images, return_tensors=None, **merged_kwargs.get("images_kwargs", {}))
-            # And text would be modified to include image tokens like self.full_image_sequence_str
-            # For now, just a pass-through to show where it fits.
-            pass  # Replace with actual image processing and text modification logic
-
-        # --- Audio Processing ---
-        audio_features = {}
+            # image_features_dict = self.image_processor(images, return_tensors=None, **merged_call_kwargs.get("images_kwargs", {}))
+            # Simplified: Actual image token replacement logic for `text` would go here.
+            # text = self._handle_image_text_replacement(text, images, image_features_dict)
+            pass
+
+
+        audio_features_dict = {}
         if audios is not None and self.audio_processor is not None:
             logger.info("Processing audio...")
-            audio_call_kwargs = merged_kwargs.get("audio_kwargs", {})
-            if sampling_rate:  # Pass sampling_rate if provided for raw arrays
-                audio_call_kwargs["sampling_rate"] = sampling_rate
+            audio_call_kwargs = merged_call_kwargs.get("audio_kwargs", {})
+            if sampling_rate: 
+                 audio_call_kwargs["sampling_rate"] = sampling_rate
+            
+            # audio_processor.__call__ returns BatchFeature, we need its .data attribute
+            audio_features_batch_feature = self.audio_processor(audios=audios, return_tensors=None, **audio_call_kwargs)
+            audio_features_dict = audio_features_batch_feature.data # Get the dict
 
-            # Get dict of numpy arrays/lists first from feature_extractor
-            audio_features_np = self.audio_processor(audios=audios, return_tensors=None, **audio_call_kwargs)
-            audio_features = audio_features_np  # Store the dict
-
-            # Modify text to include audio soft tokens
             new_text_with_audio = []
-            # audio_attention_mask from feature extractor is (B, Max_T_mel)
-            audio_sample_mel_lengths = audio_features_np["audio_attention_mask"].sum(
-                axis=-1)  # Get actual mel frames per sample
+            # audio_attention_mask shape is (B, Max_T_mel)
+            audio_sample_mel_lengths = to_py_obj(audio_features_dict["audio_attention_mask"].sum(axis=-1))
 
             for i, prompt in enumerate(text):
-                num_soft_tokens = self._compute_audio_prompt_token_count(int(audio_sample_mel_lengths[i]))
+                num_soft_tokens = self._compute_audio_prompt_token_count(audio_sample_mel_lengths[i])
                 audio_token_sequence_str = self.audio_soft_token_str * num_soft_tokens
-
+                
                 if self.audio_placeholder_token in prompt:
                     prompt = prompt.replace(self.audio_placeholder_token, audio_token_sequence_str, 1)
-                else:  # If no placeholder, append the audio tokens
-                    prompt += audio_token_sequence_str
+                else: 
+                    prompt += audio_token_sequence_str 
                 new_text_with_audio.append(prompt)
             text = new_text_with_audio
-
-        # --- Text Tokenization ---
+        
         logger.info("Tokenizing text...")
-        text_call_kwargs = merged_kwargs.get("text_kwargs", {})
-        # Ensure tokenizer gets lists/np arrays, BatchFeature will handle final tensor conversion
-        text_features = self.tokenizer(text, return_tensors=None, **text_call_kwargs)
+        text_call_kwargs = merged_call_kwargs.get("text_kwargs", {})
+        text_features_dict = self.tokenizer(text, return_tensors=None, **text_call_kwargs)
 
-        # --- Create token_type_ids ---
-        input_ids_list = text_features["input_ids"]  # Should be list of lists of token IDs
-        if not isinstance(input_ids_list[0], list):  # Handle if tokenizer returns single list for single text
-            input_ids_list = [input_ids_list]
+        input_ids_list = text_features_dict["input_ids"]
+        if not isinstance(input_ids_list, list) or not (input_ids_list and isinstance(input_ids_list[0], list)):
+             if isinstance(input_ids_list, (torch.Tensor, np.ndarray)):
+                 input_ids_list = to_py_obj(input_ids_list) # Convert tensor/np.array to list of lists
+             elif isinstance(input_ids_list, list) and (not input_ids_list or isinstance(input_ids_list[0], int)):
+                 input_ids_list = [input_ids_list] 
 
         token_type_ids_list = []
         for ids_sample in input_ids_list:
-            types = [0] * len(ids_sample)  # 0 for text
+            types = [0] * len(ids_sample) 
             for j, token_id in enumerate(ids_sample):
                 if self.image_token_id is not None and token_id == self.image_token_id:
-                    types[j] = 1  # 1 for image
-                elif token_id == self.audio_soft_token_id:  # Compare with the ID of the soft token string
-                    types[j] = 2  # 2 for audio
+                    types[j] = 1 
+                elif token_id == self.audio_soft_token_id: 
+                    types[j] = 2 
             token_type_ids_list.append(types)
-        text_features["token_type_ids"] = token_type_ids_list
-
-        # Combine all features
-        combined_features = {**text_features}
-        if image_features:  # image_features is already a dict
-            combined_features.update(image_features)
-        if audio_features:  # audio_features is already a dict from audio_processor
-            combined_features.update(audio_features)
-
+        text_features_dict["token_type_ids"] = token_type_ids_list
+        
+        combined_features = {**text_features_dict}
+        if image_features_dict: 
+            combined_features.update(image_features_dict)
+        if audio_features_dict: 
+            combined_features.update(audio_features_dict)
+            
         return BatchFeature(data=combined_features, tensor_type=final_rt)
 
     def batch_decode(self, *args, **kwargs):
@@ -532,14 +484,11 @@ class Gemma3OmniProcessor(ProcessorMixin):
 
     @property
     def model_input_names(self) -> List[str]:
-        """
-        Defines the expected inputs for the model. Combines tokenizer, image_processor, and audio_processor inputs.
-        """
         input_names = set(self.tokenizer.model_input_names + ["token_type_ids"])
         if self.image_processor is not None:
             input_names.update(self.image_processor.model_input_names)
         if self.audio_processor is not None:
-            # Gemma3AudioFeatureExtractor produces "audio_values" and "audio_attention_mask"
-            # Other keys like "audio_token_calc_sizes" are for internal use by processor, not model input
-            input_names.update(["audio_values", "audio_attention_mask"])
+            # From Gemma3AudioFeatureExtractor's output_data keys
+            input_names.update(["audio_values", "audio_attention_mask"]) 
+            # "audio_token_calc_sizes" is internal to processor, not model.
         return list(input_names)