Update processing_gemma3_omni.py

processing_gemma3_omni.py

@@ -1,5 +1,5 @@
 import re
-from typing import List, Optional, Union, Dict, Any
+from typing import List, Optional, Union, Dict, Any, Tuple
 
 import math
 import numpy as np
@@ -15,17 +15,16 @@ 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.processing_utils import ProcessorMixin, ProcessingKwargs 
+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 (as defined before)
 DEFAULT_SAMPLING_RATE = 16000
 DEFAULT_N_FFT = 512
-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_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
 DEFAULT_QFORMER_RATE = 2
@@ -48,15 +47,15 @@ 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)
-    
+
     freq_points = np.clip(freq_points, 0, sampling_rate / 2.0)
     bins = np.floor((n_fft / 2.0) * freq_points / (sampling_rate / 2.0)).astype(int)
     bins = np.clip(bins, 0, n_fft // 2)
@@ -64,43 +63,44 @@ def create_mel_filterbank(sampling_rate: int, n_fft: int, n_mels: int, fmin: flo
     filterbank = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
     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)
+            filterbank[m, left:center + 1] = (np.arange(left, center + 1) - left) / (center - left)
         if right > center:
-            filterbank[m, center:right+1] = (right - np.arange(center, right + 1)) / (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 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
-
+            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
 
+
 # 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"] 
+    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, 
+            sampling_rate: int = DEFAULT_SAMPLING_RATE,
             n_fft: int = DEFAULT_N_FFT,
-            win_length: Optional[int] = None, 
-            hop_length: Optional[int] = 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, 
+            padding_value: float = 0.0,
             **kwargs
     ):
         super().__init__(feature_size=n_mels, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
@@ -128,16 +128,16 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
     def __call__(
             self,
             audios: Union[AudioInput, List[AudioInput]],
-            sampling_rate: Optional[int] = None, 
+            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] = []
-        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:
@@ -161,11 +161,11 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
                 )
 
             processed_wav = self._preprocess_audio(current_wav_array, source_sr)
-            mel_spectrogram = self._compute_log_mel_spectrogram(processed_wav) 
-            
+            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]) 
+            actual_mel_lengths.append(feature_tensor.shape[0])
 
             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)
@@ -173,16 +173,17 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
         audio_values = pad_sequence(processed_mel_spectrograms, batch_first=True, padding_value=self.padding_value)
         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,
             "audio_attention_mask": audio_attention_mask
         }
-        
+
         if downstream_sizes_for_token_calc:
-             output_data["audio_token_calc_sizes"] = torch.stack(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:
@@ -190,22 +191,22 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
             if np.issubdtype(wav.dtype, np.integer):
                 max_val = np.iinfo(wav.dtype).max
                 wav = wav.astype(np.float32) / max_val
-            else: 
+            else:
                 wav = wav.astype(np.float32)
-        
+
         if wav.ndim > 1:
             wav = wav.mean(axis=0)
-        
+
         if source_sr != self.sampling_rate:
             gcd = math.gcd(self.sampling_rate, source_sr)
             up_factor = self.sampling_rate // gcd
             down_factor = source_sr // gcd
-            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)
-        
+            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: 
+        if norm_factor > 1e-9:
             wav = wav / norm_factor
         return wav
 
@@ -216,7 +217,8 @@ class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
 
         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(
@@ -225,21 +227,22 @@ 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)
-        powers = np.abs(stft_matrix)**2
+        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:
         compressed = math.ceil(frame_count / self.compression_rate)
         return math.ceil(compressed / self.qformer_rate)
 
-class Gemma3DummyProcessorKwargs(ProcessingKwargs, total=False): # Dummy for testing structure
+
+class Gemma3DummyProcessorKwargs(ProcessingKwargs, total=False):  # Dummy for testing structure
     images_kwargs: Dict[str, Any]
     audio_kwargs: Dict[str, Any]
     text_kwargs: Dict[str, Any]
@@ -249,24 +252,25 @@ class Gemma3DummyProcessorKwargs(ProcessingKwargs, total=False): # Dummy for tes
         "audio_kwargs": {}
     }
 
+
 class Gemma3OmniProcessor(ProcessorMixin):
     attributes = ["image_processor", "audio_processor", "tokenizer"]
     # 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
+    audio_processor_class = Gemma3AudioFeatureExtractor  # Correctly points to your custom class
+    tokenizer_class = "AutoTokenizer"  # Or the specific class string
 
     # valid_kwargs was in user's code, its role depends on ProcessorMixin internal usage
-    valid_kwargs = ["chat_template", "image_seq_length"] 
+    valid_kwargs = ["chat_template", "image_seq_length"]
 
     def __init__(
             self,
-            tokenizer, 
+            tokenizer,
             audio_processor: Optional[Union[Gemma3AudioFeatureExtractor, Dict]] = None,
-            image_processor = None, 
+            image_processor=None,
             chat_template=None,
             image_seq_length: int = 256,
-            audio_prompt_compression_rate: int = 8, 
+            audio_prompt_compression_rate: int = 8,
             audio_prompt_qformer_rate: int = 1,
             audio_prompt_feat_stride: int = 1,
             audio_placeholder_token: str = "<|audio_placeholder|>",
@@ -279,48 +283,50 @@ class Gemma3OmniProcessor(ProcessorMixin):
             audio_processor = Gemma3AudioFeatureExtractor()
         elif isinstance(audio_processor, Dict):
             audio_processor = Gemma3AudioFeatureExtractor(**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)}")
+        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
+            # 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.")
+                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
+            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)
+        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.")
-
+            raise ValueError("A tokenizer instance or identifier must be provided.")
 
         super().__init__(
             image_processor=image_processor,
             audio_processor=audio_processor,
             tokenizer=tokenizer,
             chat_template=chat_template,
-            **kwargs # Pass other kwargs to super
+            **kwargs  # Pass other kwargs to super
         )
-        
+
         self.image_seq_length = image_seq_length
-        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_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.eoi_token = getattr(self.tokenizer, "eoi_token", "")
 
         self.audio_placeholder_token = audio_placeholder_token
         self.audio_soft_token_str = audio_soft_token_str
-        
+
         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: # Check if UNK
-             logger.warning(
+        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."
             )
@@ -331,7 +337,6 @@ class Gemma3OmniProcessor(ProcessorMixin):
         self.audio_prompt_qformer_rate = audio_prompt_qformer_rate
         self.audio_prompt_feat_stride = audio_prompt_feat_stride
 
-
     def _merge_kwargs(self, KwargsClassWithDefaults, tokenizer_init_kwargs, **kwargs_passed_to_call):
         final_kwargs = {}
         # Initialize with _defaults from the Kwargs class
@@ -342,24 +347,24 @@ class Gemma3OmniProcessor(ProcessorMixin):
 
         # 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()): 
+            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__
         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()
+                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"] = {}  # 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:
@@ -371,13 +376,13 @@ class Gemma3OmniProcessor(ProcessorMixin):
     def __call__(
             self,
             text: Union[str, List[str]] = None,
-            images: Optional[Any] = None, 
+            images: Optional[Any] = None,
             audios: Optional[Union[AudioInput, List[AudioInput]]] = None,
-            sampling_rate: Optional[int] = None, 
-            return_tensors: Optional[Union[str, TensorType]] = None, 
-            **kwargs: Any 
+            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`.")
 
@@ -385,27 +390,27 @@ class Gemma3OmniProcessor(ProcessorMixin):
         # 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
+            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__
+
+        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
+        else:  # If passed directly, remove from text_kwargs to avoid conflict
             merged_call_kwargs.get("text_kwargs", {}).pop("return_tensors", None)
 
-
         if text is None:
             num_samples = 0
             if images is not None:
-                _images_list = images if isinstance(images, list) and (not images or not isinstance(images[0], (int, float))) else [images]
+                _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:
                 _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)):
@@ -419,17 +424,16 @@ class Gemma3OmniProcessor(ProcessorMixin):
             # 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_call_kwargs.get("audio_kwargs", {})
-            if sampling_rate: 
-                 audio_call_kwargs["sampling_rate"] = sampling_rate
-            
+            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
+            audio_features_dict = audio_features_batch_feature.data  # Get the dict
 
             new_text_with_audio = []
             # audio_attention_mask shape is (B, Max_T_mel)
@@ -438,42 +442,42 @@ class Gemma3OmniProcessor(ProcessorMixin):
             for i, prompt in enumerate(text):
                 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: 
-                    prompt += audio_token_sequence_str 
+                else:
+                    prompt += audio_token_sequence_str
                 new_text_with_audio.append(prompt)
             text = new_text_with_audio
-        
+
         logger.info("Tokenizing text...")
         text_call_kwargs = merged_call_kwargs.get("text_kwargs", {})
         text_features_dict = self.tokenizer(text, return_tensors=None, **text_call_kwargs)
 
         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] 
+            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) 
+            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 
-                elif token_id == self.audio_soft_token_id: 
-                    types[j] = 2 
+                    types[j] = 1
+                elif token_id == self.audio_soft_token_id:
+                    types[j] = 2
             token_type_ids_list.append(types)
         text_features_dict["token_type_ids"] = token_type_ids_list
-        
+
         combined_features = {**text_features_dict}
-        if image_features_dict: 
+        if image_features_dict:
             combined_features.update(image_features_dict)
-        if audio_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):
@@ -489,6 +493,6 @@ class Gemma3OmniProcessor(ProcessorMixin):
             input_names.update(self.image_processor.model_input_names)
         if self.audio_processor is not None:
             # From Gemma3AudioFeatureExtractor's output_data keys
-            input_names.update(["audio_values", "audio_attention_mask"]) 
+            input_names.update(["audio_values", "audio_attention_mask"])
             # "audio_token_calc_sizes" is internal to processor, not model.
         return list(input_names)