Upload 12 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1 @@
+{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' %}{% set loop_messages = messages[1:] %}{% else %}{% set first_user_prefix = '' %}{% set loop_messages = messages %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set role = 'model' if message['role'] == 'assistant' else message['role'] %}{{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else '') }}{% if role == 'model' and message.get('metadata') %}{% if message['metadata']['type'] == 'think' %}<think>{% if message['metadata'].get('range') %}<range>{{ message['metadata']['range'] }}</range>{% endif %}{% if message['metadata'].get('content') %}{{ message['metadata']['content'] | trim }}{% endif %}</think>{% elif message['metadata']['type'] == 'direct' %}<direct>{% endif %}{% if message['metadata'].get('function') %}<function>{{ message['metadata']['function'] | join(',') }}</function>{% endif %}{% endif %}{% if message['content'] is string %}{{ message['content'] | trim }}{% elif message['content'] is iterable %}{% for item in message['content'] %}{{ '<start_of_image>' if item['type']=='image' else '<start_of_audio>' if item['type']=='audio' else item['text'] | trim if item['type']=='text' else '' }}{% endfor %}{% else %}{{ raise_exception('Invalid content type') }}{% endif %}{{ '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}

config.json

@@ -2,7 +2,13 @@
   "architectures": [
     "Gemma3OmniForConditionalGeneration"
   ],
-  "audio_token_index": 262151,
+  "auto_map": {
+      "AutoProcessor": "processing_gemma3_omni.Gemma3OmniProcessor",
+      "AutoFeatureExtractor": "processing_gemma3_omni.Gemma3AudioFeatureExtractor",
+      "AutoModel": "modeling_gemma_3_omni.Gemma3OmniForConditionalGeneration",
+      "AutoModelForCausalLM":  "modeling_gemma3_omni.Gemma3OmniForConditionalGeneration",
+      "AutoConfig": "configuration_gemma3_omni.Gemma3OmniConfig"
+  },
   "boi_token_index": 255999,
   "eoi_token_index": 256000,
   "eos_token_id": [
@@ -122,4 +128,4 @@
     "torch_dtype": "float32",
     "vision_use_head": false
   }
-}
+}

configuration_gemma3_omni.py

@@ -0,0 +1,55 @@
+from typing import Optional, Union, Dict, Any
+
+from transformers import Gemma3TextConfig, SiglipVisionConfig, PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class Gemma3OmniConfig(PretrainedConfig):
+    model_type = "gemma3omni"
+    attribute_map = {
+        "image_token_id": "image_token_index",
+        "audio_token_id": "audio_token_index",
+        "boi_token_id": "boi_token_index",
+        "eoi_token_id": "eoi_token_index",
+    }
+    sub_configs = {
+        "text_config": Gemma3TextConfig,
+        "vision_config": SiglipVisionConfig,
+    }
+
+    def __init__(
+            self,
+            text_config: Optional[Union[Gemma3TextConfig, Dict[str, Any]]] = None,
+            vision_config: Optional[Union[SiglipVisionConfig, Dict[str, Any]]] = None,
+            mm_tokens_per_image: int = 256,
+            boi_token_index: int = 255_999,
+            eoi_token_index: int = 256_000,
+            image_token_index: int = 262_144,
+            audio_token_index: int = 262_151,
+            initializer_range: float = 0.02,
+            **kwargs,
+    ):
+        if text_config is None:
+            text_config = Gemma3TextConfig()
+            logger.info("text_config is None, using default Gemma3TextConfig text config.")
+        elif isinstance(text_config, dict):
+            text_config = Gemma3TextConfig(**text_config)
+
+        if isinstance(vision_config, dict):
+            vision_config = SiglipVisionConfig(**vision_config)
+        elif vision_config is None:
+            vision_config = SiglipVisionConfig()
+            logger.info("vision_config is None, using default SiglipVisionConfig vision config.")
+
+        self.text_config = text_config
+        self.vision_config = vision_config
+        self.mm_tokens_per_image = mm_tokens_per_image
+        self.boi_token_index = boi_token_index
+        self.eoi_token_index = eoi_token_index
+        self.image_token_index = image_token_index
+        self.audio_token_index = audio_token_index
+        self.initializer_range = initializer_range
+
+        super().__init__(**kwargs)

modeling_gemma3_omni.py

@@ -0,0 +1,461 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+from typing import List, Optional, Tuple, Union, Callable
+
+from transformers import (
+    AutoModel,
+    Cache,
+    PreTrainedModel,
+    PretrainedConfig, )
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask
+from transformers.models.gemma3.modeling_gemma3 import (
+    Gemma3CausalLMOutputWithPast,
+    Gemma3RMSNorm, Gemma3PreTrainedModel, Gemma3ModelOutputWithPast,
+)
+from transformers.utils import is_torchdynamo_compiling, logging, is_torch_flex_attn_available
+
+try:
+    from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss
+except:
+    LigerFusedLinearCrossEntropyLoss = None
+
+from .configuration_gemma3_omni import Gemma3OmniConfig
+from .speech_conformer_encoder import ConformerEncoder
+
+logger = logging.get_logger(__name__)
+
+if is_torch_flex_attn_available():
+    from torch.nn.attention.flex_attention import BlockMask
+
+
+class Gemma3AudioProjectorConfig(PretrainedConfig):
+    model_type = "gemma3_audio"
+
+    def __init__(
+            self,
+            hidden_size: int = 1024,
+            num_hidden_layers: int = 24,
+            sample_rate: int = 16_000,
+            n_mels: int = 80,
+            image_token_index: int = 0,  # This seems unused for audio projector, maybe a copy-paste?
+            # Added Mel spectrogram specific parameters
+            n_fft: int = 400,  # Typical for 25ms window at 16kHz
+            hop_length: int = 160,  # Typical for 10ms hop at 16kHz
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.sample_rate = sample_rate
+        self.n_mels = n_mels
+        self.image_token_index = image_token_index
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+
+
+import torch
+from torch import nn
+
+
+class LayerWiseWeightedSum(nn.Module):
+    def __init__(self, num_layers: int, learnable: bool = True):
+        super().__init__()
+        self.num_layers = num_layers
+        if learnable:
+            self.scalar = nn.Parameter(torch.zeros(num_layers))
+        else:
+            self.register_buffer("scalar", torch.zeros(num_layers))
+
+    def forward(self, hidden_states: list[torch.Tensor]) -> torch.Tensor:
+        assert len(hidden_states) == self.num_layers
+        norm_w = torch.softmax(self.scalar, dim=0).view(-1, 1, 1, 1)
+        stacked = torch.stack(hidden_states, dim=0)
+        return (norm_w * stacked).sum(dim=0)
+
+
+class Gemma3AudioProjector(PreTrainedModel):
+    """Conformer-based audio encoder → project to LM hidden-dim."""
+
+    config_class = Gemma3AudioProjectorConfig
+    base_model_prefix = "audio_projector"
+
+    def __init__(self, config: Gemma3AudioProjectorConfig):
+        super().__init__(config)
+        encoder_config = {
+            "activation": "swish",
+            "activation_checkpointing": "",
+            "attention_dim": 1024,
+            "attention_heads": 16,
+            "batch_norm": False,
+            "bias_in_glu": True,
+            "causal": True,
+            "chunk_size": -1,
+            "conv_activation": "swish",
+            "conv_glu_type": "swish",
+            "depthwise_multiplier": 1,
+            "depthwise_seperable_out_channel": 1024,
+            "dropout_rate": 0.0,
+            "encoder_embedding_config": {
+                "input_size": config.n_mels  # This is feat_in for NemoConvSubsampling
+            },
+            "ext_pw_kernel_size": 1,
+            "ext_pw_out_channel": 1024,
+            "input_layer": "nemo_conv",
+            "input_size": config.n_mels,  # Also feat_in for NemoConvSubsampling, consistency
+            "kernel_size": 3,
+            "left_chunk": 18,
+            "linear_units": 1536,
+            "nemo_conv_settings": {
+                "conv_channels": 1024,
+            },
+            "num_blocks": 24,
+            "relative_attention_bias_args": {
+                "t5_bias_max_distance": 500,
+                "type": "t5"
+            },
+            "time_reduction": 8
+        }
+        self.encoder = ConformerEncoder(**encoder_config)
+        self.layer_weighter = LayerWiseWeightedSum(
+            num_layers=encoder_config["num_blocks"]
+        )
+        self.proj = nn.Linear(encoder_config['attention_dim'], config.hidden_size, bias=False)
+
+    def forward(self, mel: torch.Tensor, mel_mask: torch.Tensor):
+        mel = mel.squeeze(1)  # (B, T, 80)
+        mel_mask = mel_mask.squeeze(1)  # (B, L)
+
+        if mel_mask.size(1) != mel.size(1):
+            mel_mask = mel_mask[..., : mel.size(1)]
+
+        _, out_mask, hidden_list = self.encoder(mel, mel_mask)
+        hidden_sum = self.layer_weighter(hidden_list)
+        hidden = self.proj(hidden_sum)
+        return hidden, out_mask
+
+
+class Gemma3VisionProjector(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.mm_input_projection_weight = nn.Parameter(
+            torch.zeros(config.vision_config.hidden_size, config.text_config.hidden_size)
+        )
+        self.mm_soft_emb_norm = Gemma3RMSNorm(
+            config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps
+        )
+        self.patches_per_image = config.vision_config.image_size // config.vision_config.patch_size
+        self.tokens_per_side = int(config.mm_tokens_per_image ** 0.5)
+        self.kernel_size = self.patches_per_image // self.tokens_per_side
+        self.avg_pool = nn.AvgPool2d(kernel_size=self.kernel_size, stride=self.kernel_size)
+
+    def forward(self, vision_outputs: torch.Tensor):
+        b, _, seq_len = vision_outputs.shape
+        x = vision_outputs.transpose(1, 2).reshape(
+            b, seq_len, self.patches_per_image, self.patches_per_image
+        )
+        x = self.avg_pool(x).flatten(2).transpose(1, 2)
+        x = self.mm_soft_emb_norm(x)
+        return torch.matmul(x, self.mm_input_projection_weight).type_as(vision_outputs)
+
+
+def token_type_ids_mask_function(token_type_ids: Optional[torch.Tensor]) -> Optional[Callable]:
+    """
+    This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,
+    not start and end indices.
+    """
+    # Do not return an additional mask in this case
+    if token_type_ids is None:
+        return None
+
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        # If it's 1, we need to unmask it
+        return token_type_ids[batch_idx, kv_idx] == 1
+
+    return inner_mask
+
+
+class Gemma3OmniModel(Gemma3PreTrainedModel):
+    config_class = Gemma3OmniConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.vision_tower = AutoModel.from_config(config=config.vision_config)
+        self.multi_modal_projector = Gemma3VisionProjector(config)
+        self.audio_projector = Gemma3AudioProjector(
+            Gemma3AudioProjectorConfig(hidden_size=config.text_config.hidden_size)
+        )
+        self.vocab_size = config.text_config.vocab_size
+
+        language_model = AutoModel.from_config(config=config.text_config)
+        self.language_model = language_model
+
+        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.language_model.embed_tokens
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            pixel_values: torch.FloatTensor = None,
+            input_audio_embeds: Optional[torch.FloatTensor] = None,
+            audio_attention_mask: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None,
+            token_type_ids: Optional[torch.LongTensor] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            **lm_kwargs,
+    ) -> Union[Tuple, Gemma3ModelOutputWithPast]:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            print("input_ids:", input_ids, "inputs_embeds:", inputs_embeds)
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Replace image id woth PAD if the image token if OOV, to avoid index-errors
+        if input_ids is not None and self.config.image_token_id >= self.vocab_size:
+            special_image_mask = input_ids == self.config.image_token_id
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[special_image_mask] = 0
+        else:
+            llm_input_ids = input_ids
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if pixel_values is not None and past_key_values is None:
+            image_features = self.get_image_features(pixel_values)
+
+            if input_ids is None:
+                special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                )
+            else:
+                special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+                special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+
+            if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
+                image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
+                raise ValueError(
+                    f"Number of images does not match number of special image tokens in the input text. "
+                    f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
+                    "tokens from image embeddings."
+                )
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+        if input_audio_embeds is not None and past_key_values is None:
+            audio_features, audio_feat_mask = self.audio_projector(
+                input_audio_embeds, audio_attention_mask
+            )
+            if input_ids is None:
+                special_audio_mask = (
+                        inputs_embeds
+                        == self.get_input_embeddings()(
+                    torch.tensor(
+                        self.config.audio_token_index,
+                        dtype=torch.long,
+                        device=inputs_embeds.device,
+                    )
+                )
+                )
+            else:
+                special_audio_mask = (
+                        input_ids == self.config.audio_token_index
+                ).unsqueeze(-1)  # [B, L, 1]
+                special_audio_mask = special_audio_mask.expand_as(inputs_embeds).to(
+                    inputs_embeds.device
+                )
+            if (
+                    not is_torchdynamo_compiling()
+                    and inputs_embeds[special_audio_mask].numel() != audio_features.numel()
+            ):
+                audio_tokens_in_text = special_audio_mask.sum(dim=1).sum(dim=0)[0]
+                raise ValueError(
+                    f"Number of audio tokens in the text ({audio_tokens_in_text}) "
+                    f"≠ number of tokens from audio embeddings "
+                    f"({audio_features.shape[0] * audio_features.shape[1]})."
+                )
+            audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            audio_features = audio_features.reshape(-1)
+            inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config.get_text_config(),
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+            }
+            if token_type_ids is not None and inputs_embeds.shape[1] != 1:
+                mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
+                    token_type_ids.to(cache_position.device)
+                )
+
+            # Create the masks
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+            }
+
+        outputs = self.language_model(
+            attention_mask=causal_mask_mapping,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **lm_kwargs,
+        )
+
+        return Gemma3ModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values if use_cache else None,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features if pixel_values is not None else None,
+        )
+
+
+class Gemma3OmniForConditionalGeneration(Gemma3PreTrainedModel, GenerationMixin):
+    config_class = Gemma3OmniConfig
+    """Gemma-3 Omni：vision + audio + text causal LM."""
+    _checkpoint_conversion_mapping = {
+        "^language_model.model": "model.language_model",
+        "^vision_tower": "model.vision_tower",
+        "^multi_modal_projector": "model.multi_modal_projector",
+        "^language_model.lm_head": "lm_head",
+    }
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Gemma3OmniModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.language_model.embed_tokens
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            pixel_values: torch.FloatTensor = None,
+            input_audio_embeds: Optional[torch.FloatTensor] = None,
+            audio_attention_mask: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None,
+            token_type_ids: Optional[torch.LongTensor] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            logits_to_keep: Union[int, torch.Tensor] = 0,
+            **lm_kwargs,
+    ) -> Union[Tuple, Gemma3CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            input_audio_embeds=input_audio_embeds,
+            audio_attention_mask=audio_attention_mask,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            labels=labels,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **lm_kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            shift_logits = logits[..., :-1, :]
+            shift_labels = labels[..., 1:]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1]:].to(logits.device)
+                shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+
+            if LigerFusedLinearCrossEntropyLoss is not None:
+                loss_fct = LigerFusedLinearCrossEntropyLoss()
+            else:
+                loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(flat_logits, flat_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return Gemma3CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+        )
+
+
+__all__ = [
+    "Gemma3AudioProjectorConfig",
+    "Gemma3AudioProjector",
+    "Gemma3VisionProjector",
+    "Gemma3OmniModel",
+    "Gemma3OmniForConditionalGeneration",
+]

preprocessor_config.json

@@ -0,0 +1,46 @@
+{
+  "auto_map": {
+      "AutoProcessor": "processing_gemma3_omni.Gemma3OmniProcessor",
+      "AutoFeatureExtractor": "processing_gemma3_omni.Gemma3AudioFeatureExtractor"
+  },
+  "do_convert_rgb": null,
+  "do_normalize": true,
+  "do_pan_and_scan": null,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Gemma3ImageProcessor",
+  "processor_class": "Gemma3Processor",
+  "image_seq_length": 256,
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "pan_and_scan_max_num_crops": null,
+  "pan_and_scan_min_crop_size": null,
+  "pan_and_scan_min_ratio_to_activate": null,
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 896,
+    "width": 896
+  },
+  "compression_rate": 4,
+  "feat_stride": 4,
+  "feature_extractor_type": "Gemma3AudioFeatureExtractor",
+  "feature_size": 80,
+  "hop_length": 160,
+  "n_fft": 512,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "processor_class": "Gemma3OmniProcessor",
+  "qformer_rate": 2,
+  "return_attention_mask": true,
+  "sampling_rate": 16000,
+  "win_length": 400
+}

processing_gemma3_omni.py

@@ -0,0 +1,491 @@
+from typing import List, Optional, Union, Dict, Any, Tuple
+
+import numpy as np
+import scipy.signal
+import torch
+from torch.nn.utils.rnn import pad_sequence
+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
+from transformers.processing_utils import ProcessorMixin
+from transformers.utils import TensorType, logging
+
+DEFAULT_SPECIAL_TOKENS = {
+    "bos_token": "<bos>",
+    "eos_token": "<eos>",
+    "pad_token": "<pad>",
+    "unk_token": "<unk>",
+    "boi_token": "<start_of_image>",
+    "eoi_token": "<end_of_image>",
+    "image_token": "<image_soft_token>",
+    "boa_token": "<start_of_audio>",
+    "eoa_token": "<end_of_audio>",
+    "audio_token": "<audio_soft_token>",
+}
+DEFAULT_SAMPLING_RATE = 16000
+DEFAULT_N_FFT = 512
+DEFAULT_WIN_LENGTH = 400
+DEFAULT_HOP_LENGTH = 160
+DEFAULT_N_MELS = 80
+DEFAULT_COMPRESSION_RATE = 4
+DEFAULT_QFORMER_RATE = 8
+DEFAULT_FEAT_STRIDE = 4
+DEFAULT_IMAGE_SEQ_LENGTH = 256
+DEFAULT_MAX_LENGTH = 16384
+
+logger = logging.get_logger(__name__)
+
+
+def compute_audio_token_count(
+        mel_frame_count: int,
+        *,
+        feat_stride: int = DEFAULT_FEAT_STRIDE,
+        compression_rate: int = DEFAULT_COMPRESSION_RATE,
+        qformer_rate: int = DEFAULT_QFORMER_RATE,
+) -> int:
+    audio_frames = mel_frame_count * feat_stride
+    audio_frames = (audio_frames + compression_rate - 1) // compression_rate
+    audio_frames = (audio_frames + qformer_rate - 1) // qformer_rate
+    return audio_frames
+
+
+def speechlib_mel(sample_rate, n_fft, n_mels, fmin=None, fmax=None):
+    bank_width = int(n_fft // 2 + 1)
+    if fmax is None:
+        fmax = sample_rate / 2
+    if fmin is None:
+        fmin = 0
+
+    def mel(f):
+        return 1127.0 * np.log(1.0 + f / 700.0)
+
+    def bin2mel(fft_bin):
+        return 1127.0 * np.log(1.0 + fft_bin * sample_rate / (n_fft * 700.0))
+
+    def f2bin(f):
+        return int((f * n_fft / sample_rate) + 0.5)
+
+    klo = f2bin(fmin) + 1
+    khi = f2bin(fmax)
+    khi = max(khi, klo)
+    mlo = mel(fmin)
+    mhi = mel(fmax)
+    m_centers = np.linspace(mlo, mhi, n_mels + 2)
+    ms = (mhi - mlo) / (n_mels + 1)
+    matrix = np.zeros((n_mels, bank_width), dtype=np.float32)
+    for m in range(n_mels):
+        left = m_centers[m]
+        center = m_centers[m + 1]
+        right = m_centers[m + 2]
+        for fft_bin in range(klo, khi):
+            mbin = bin2mel(fft_bin)
+            if left < mbin < right:
+                matrix[m, fft_bin] = 1.0 - abs(center - mbin) / ms
+    return matrix
+
+
+class Gemma3AudioFeatureExtractor(SequenceFeatureExtractor):
+    model_input_names = ["input_audio_embeds", "audio_embed_sizes", "audio_attention_mask"]
+
+    def __init__(
+            self,
+            audio_compression_rate: int = DEFAULT_COMPRESSION_RATE,
+            audio_downsample_rate: int = DEFAULT_QFORMER_RATE,
+            audio_feat_stride: int = DEFAULT_FEAT_STRIDE,
+            feature_size: int = DEFAULT_N_MELS,
+            sampling_rate: int = DEFAULT_SAMPLING_RATE,
+            padding_value: float = 0.0,
+            eightk_method: str = "fillzero",
+            **kwargs,
+    ):
+        super().__init__(
+            feature_size=kwargs.pop("feature_size", feature_size),
+            sampling_rate=kwargs.pop("sampling_rate", sampling_rate),
+            padding_value=kwargs.pop("padding_value", padding_value),
+            **kwargs,
+        )
+        self.compression_rate = audio_compression_rate
+        self.qformer_compression_rate = audio_downsample_rate
+        self.feat_stride = audio_feat_stride
+        self._eightk_method = eightk_method
+        self._mel = speechlib_mel(16000, 512, 80, fmin=None, fmax=7690).T
+        self._hamming400 = np.hamming(400)
+        self._hamming200 = np.hamming(200)
+
+    def __call__(
+            self,
+            audios: List[Tuple[np.ndarray, int]],
+            return_tensors: Optional[Union[str, TensorType]] = None,
+    ):
+        returned_input_audio_embeds = []
+        returned_audio_embed_sizes = []
+        audio_frames_list = []
+        for audio_data, sample_rate in audios:
+            if isinstance(audio_data, list):
+                audio_data = np.array(audio_data, dtype=np.float32)
+            if not isinstance(audio_data, np.ndarray):
+                raise TypeError(f"Waveform data must be a numpy array, got {type(audio_data)}")
+            audio_embeds_np = self._extract_features(audio_data, sample_rate)
+            num_mel_frames = audio_embeds_np.shape[0]
+            current_audio_frames = num_mel_frames * self.feat_stride
+            audio_embed_size = self._compute_audio_embed_size(current_audio_frames)
+            returned_input_audio_embeds.append(torch.from_numpy(audio_embeds_np))
+            returned_audio_embed_sizes.append(torch.tensor(audio_embed_size).long())
+            audio_frames_list.append(current_audio_frames)
+        padded_input_audio_embeds = pad_sequence(
+            returned_input_audio_embeds, batch_first=True, padding_value=self.padding_value
+        )
+        stacked_audio_embed_sizes = torch.stack(returned_audio_embed_sizes, dim=0)
+        tensor_audio_frames = torch.tensor(audio_frames_list, dtype=torch.long)
+        max_audio_frames = tensor_audio_frames.max().item() if tensor_audio_frames.numel() > 0 else 0
+        if max_audio_frames > 0 and len(audios) > 1:
+            audio_attention_mask = (
+                    torch.arange(0, max_audio_frames, device=tensor_audio_frames.device).unsqueeze(0)
+                    < tensor_audio_frames.unsqueeze(1)
+            )
+        elif max_audio_frames > 0:
+            audio_attention_mask = torch.ones(1, max_audio_frames, dtype=torch.bool, device=tensor_audio_frames.device)
+        else:
+            audio_attention_mask = None
+        data = {
+            "input_audio_embeds": padded_input_audio_embeds,
+            "audio_embed_sizes": stacked_audio_embed_sizes,
+        }
+        if audio_attention_mask is not None:
+            data["audio_attention_mask"] = audio_attention_mask
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def _extract_spectrogram(self, wav: np.ndarray, fs: int) -> np.ndarray:
+        if wav.ndim > 1:
+            wav = np.squeeze(wav)
+        if len(wav.shape) == 2:
+            wav = wav.mean(axis=1).astype(np.float32)
+        wav = wav.astype(np.float32)
+        current_fs = fs
+        if current_fs > self.sampling_rate:
+            wav = scipy.signal.resample_poly(wav, self.sampling_rate, current_fs)
+            current_fs = self.sampling_rate
+        elif 8000 < current_fs < self.sampling_rate:
+            wav = scipy.signal.resample_poly(wav, 8000, current_fs)
+            current_fs = 8000
+        elif current_fs < 8000 and current_fs > 0:
+            wav = scipy.signal.resample_poly(wav, 8000, current_fs)
+            current_fs = 8000
+        elif current_fs <= 0:
+            raise RuntimeError(f"Unsupported sample rate {current_fs}")
+        if current_fs == 8000 and self._eightk_method == "resample":
+            wav = scipy.signal.resample_poly(wav, self.sampling_rate, 8000)
+            current_fs = self.sampling_rate
+        elif current_fs != self.sampling_rate:
+            raise RuntimeError(
+                f"Audio sample rate {current_fs} not supported. Expected {self.sampling_rate} or 8000 for 8k methods.")
+        preemphasis_coeff = 0.97
+        if current_fs == 8000:
+            n_fft, win_length, hop_length, fft_window = 256, 200, 80, self._hamming200
+        else:
+            n_fft, win_length, hop_length, fft_window = 512, 400, 160, self._hamming400
+        if len(wav) < win_length:
+            wav = np.pad(wav, (0, win_length - len(wav)), 'constant', constant_values=(0.0,))
+        num_frames = (wav.shape[0] - win_length) // hop_length + 1
+        if num_frames <= 0:
+            return np.zeros((0, n_fft // 2 + 1), dtype=np.float32)
+        y_frames = np.array(
+            [wav[i * hop_length: i * hop_length + win_length] for i in range(num_frames)],
+            dtype=np.float32,
+        )
+        _y_frames_rolled = np.roll(y_frames, 1, axis=1)
+        _y_frames_rolled[:, 0] = _y_frames_rolled[:, 1]
+        y_frames_preemphasized = (y_frames - preemphasis_coeff * _y_frames_rolled) * 32768.0
+        S = np.fft.rfft(fft_window * y_frames_preemphasized, n=n_fft, axis=1).astype(np.complex64)
+        if current_fs == 8000 and self._eightk_method == "fillzero":
+            target_bins = (512 // 2) + 1
+            S_core = S[:, :-1]
+            padarray = np.zeros((S_core.shape[0], target_bins - S_core.shape[1]), dtype=S.dtype)
+            S = np.concatenate((S_core, padarray), axis=1)
+        spec = np.abs(S).astype(np.float32)
+        return spec
+
+    def _extract_features(self, wav: np.ndarray, fs: int) -> np.ndarray:
+        spec = self._extract_spectrogram(wav, fs)
+        if spec.shape[0] == 0:
+            return np.zeros((0, self.feature_size), dtype=np.float32)
+        spec_power = spec ** 2
+        fbank_power = np.clip(spec_power.dot(self._mel), 1.0, None)
+        log_fbank = np.log(fbank_power).astype(np.float32)
+        return log_fbank
+
+    def _compute_audio_embed_size(self, audio_frames: int) -> int:
+        integer = audio_frames // self.compression_rate
+        remainder = audio_frames % self.compression_rate
+        result = integer if remainder == 0 else integer + 1
+        integer = result // self.qformer_compression_rate
+        remainder = result % self.qformer_compression_rate
+        result = integer if remainder == 0 else integer + 1
+        return result
+
+
+class Gemma3OmniProcessor(ProcessorMixin):
+    attributes = ["image_processor", "audio_processor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    audio_processor_class = "AutoFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+            self,
+            image_processor=None,
+            audio_processor=None,
+            tokenizer=None,
+            special_tokens: Optional[Dict[str, str]] = None,
+            image_seq_length: int = DEFAULT_IMAGE_SEQ_LENGTH,
+            prompt_audio_compression_rate: int = DEFAULT_COMPRESSION_RATE,
+            prompt_audio_qformer_rate: int = DEFAULT_QFORMER_RATE,
+            audio_placeholder_token: str = "<|audio_placeholder|>",
+            **kwargs,
+    ):
+        super().__init__(
+            image_processor=image_processor,
+            audio_processor=audio_processor,
+            tokenizer=tokenizer,
+            **kwargs,
+        )
+        self.special_tokens = dict(DEFAULT_SPECIAL_TOKENS)
+        if special_tokens is not None:
+            self.special_tokens.update(special_tokens)
+        if tokenizer is not None:
+            for key in self.special_tokens:
+                val = getattr(tokenizer, key, None)
+                if isinstance(val, str):
+                    self.special_tokens[key] = val
+        self.image_token = self.special_tokens["image_token"]
+        self.audio_token = self.special_tokens["audio_token"]
+        self.boi_token = self.special_tokens["boi_token"]
+        self.eoi_token = self.special_tokens["eoi_token"]
+        self.boa_token = self.special_tokens["boa_token"]
+        self.eoa_token = self.special_tokens["eoa_token"]
+        self.image_seq_length = image_seq_length
+        self.full_image_sequence = f"{self.boi_token}{''.join([self.image_token] * self.image_seq_length)}{self.eoi_token}"
+        self.prompt_audio_compression_rate = prompt_audio_compression_rate
+        self.prompt_audio_qformer_rate = prompt_audio_qformer_rate
+        self.audio_placeholder_token = audio_placeholder_token
+        if self.tokenizer is not None:
+            self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
+            self.audio_token_id = self.tokenizer.convert_tokens_to_ids(self.audio_token)
+        else:
+            self.image_token_id = None
+            self.audio_token_id = None
+
+    def compute_audio_token_count(self, mel_frame_count: int) -> int:
+        stride = getattr(self.audio_processor, "feat_stride", DEFAULT_FEAT_STRIDE)
+        return compute_audio_token_count(
+            mel_frame_count,
+            feat_stride=stride,
+            compression_rate=self.prompt_audio_compression_rate,
+            qformer_rate=self.prompt_audio_qformer_rate,
+        )
+
+    def apply_chat_template(
+            self,
+            messages,
+            add_generation_prompt: bool = True,
+            tokenize: bool = False,
+            **kwargs
+    ) -> Union[str, Dict[str, Any]]:
+        prompt = ""
+        if isinstance(messages, dict) and "messages" in messages:
+            if "audios" in messages:
+                audios = messages["audios"]
+            if "audio" in messages:
+                audios = [messages["audio"]]
+            if "images" in messages:
+                images = messages["images"]
+            if "image" in messages:
+                images = [messages["image"]]
+            messages = messages["messages"]
+
+        for msg in messages:
+            role = msg.get("role", "")
+            prompt += f"<start_of_turn>{role}\n"
+            contents = msg.get("content", [])
+            if not isinstance(contents, list):
+                contents = [contents]
+
+            for c in contents:
+                if isinstance(c, dict):
+                    ctype = c.get("type")
+                    if ctype == "image":
+                        idx = c.get("index")
+                        img_data = None
+                        if idx is not None and isinstance(idx, int):
+                            img_data = images[idx]
+                        elif "image" in c:
+                            img_data = c["image"]
+                        if img_data is None:
+                            logger.warning("No image data found for image content: %s", c)
+                        prompt += self.full_image_sequence
+                        continue
+
+                    if ctype == "audio":
+                        idx = c.get("index")
+                        aud_data = None
+                        if idx is not None and isinstance(idx, int):
+                            aud_data = audios[idx]["array"]
+                            sr = audios[idx].get("sampling_rate",
+                                                 self.audio_processor.sampling_rate if self.audio_processor else DEFAULT_SAMPLING_RATE)
+                        elif "audio" in c:
+                            aud_data = c["audio"]
+                            sr = c.get("sampling_rate",
+                                       self.audio_processor.sampling_rate if self.audio_processor else DEFAULT_SAMPLING_RATE)
+                        if aud_data is None:
+                            logger.warning("No audio data found for audio content: %s", c)
+
+                        n_audio_tokens = 0
+                        if self.audio_processor:
+                            features = self.audio_processor(audios=[(aud_data, sr)], return_tensors=None)
+                            mel_frame_count = features["input_audio_embeds"].shape[1]
+                            n_audio_tokens = self.compute_audio_token_count(mel_frame_count)
+                        prompt += (
+                                self.boa_token +
+                                (self.audio_token * n_audio_tokens) +
+                                self.eoa_token
+                        )
+                        continue
+
+                    if ctype == "text" and "text" in c:
+                        prompt += str(c["text"])
+                        continue
+                    continue
+
+                elif isinstance(c, str):
+                    prompt += c
+                    continue
+                else:
+                    logger.warning("Unknown content type in message: %s", c)
+                    continue
+
+            prompt += "<end_of_turn>\n"
+
+        if add_generation_prompt:
+            prompt += "<start_of_turn>model\n"
+
+        if tokenize and self.tokenizer is not None:
+            safe_kwargs = {}
+            allowed_keys = {"return_tensors", "padding", "truncation", "max_length", "add_special_tokens"}
+            for k, v in kwargs.items():
+                if k in allowed_keys:
+                    safe_kwargs[k] = v
+            return self.tokenizer(prompt, **safe_kwargs)
+
+        return prompt
+
+    def __call__(
+            self,
+            text: Optional[Union[str, List[str]]] = None,
+            images: Optional[Union[Any, List[Any]]] = None,
+            audios: Optional[Union[Tuple[np.ndarray, int], List[Tuple[np.ndarray, int]]]] = None,
+            messages: Optional[List[Dict]] = None,
+            add_generation_prompt: bool = True,
+            return_tensors: Optional[Union[str, TensorType]] = "pt",
+            device: Optional[str] = None,
+            **kwargs
+    ) -> Dict[str, Any]:
+        if messages is not None:
+            if isinstance(messages, dict):
+                messages = [messages]
+            prompt = self.apply_chat_template(
+                messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=False,
+            )
+            audio_inputs = []
+            for msg in messages:
+                contents = msg.get("content", [])
+                if not isinstance(contents, list):
+                    contents = [contents]
+                for c in contents:
+                    if isinstance(c, dict) and c.get("type") == "audio":
+                        arr = c["audio"]
+                        sr = c.get("sampling_rate",
+                                   self.audio_processor.sampling_rate if self.audio_processor else 16000)
+                        audio_inputs.append((arr, sr))
+            audio_features = {}
+            if audio_inputs and self.audio_processor is not None:
+                audio_features = self.audio_processor(audios=audio_inputs, return_tensors=return_tensors)
+            text_features = self.tokenizer(prompt, return_tensors=return_tensors, padding=True, truncation=True,
+                                           max_length=DEFAULT_MAX_LENGTH)
+            inputs = {**text_features, **audio_features}
+        else:
+            if text is None and images is None and audios is None:
+                raise ValueError("At least one of text/images/audios/messages must be provided.")
+            num_samples = 1
+            if isinstance(text, list):
+                num_samples = len(text)
+            elif images is not None and isinstance(images, list):
+                num_samples = len(images)
+            elif audios is not None and isinstance(audios, list):
+                num_samples = len(audios)
+            image_features = {}
+            if images is not None and self.image_processor is not None:
+                batched_images = make_nested_list_of_images(images)
+                img_out = self.image_processor(batched_images, return_tensors=None)
+                image_features = img_out.data if isinstance(img_out, BatchFeature) else img_out
+            audio_features = {}
+            audio_token_counts = None
+            if audios is not None and self.audio_processor is not None:
+                audio_out = self.audio_processor(audios=audios, return_tensors=None)
+                audio_features = audio_out.data
+                att_mask = audio_features[self.audio_processor.model_input_names[2]]
+                if isinstance(att_mask, torch.Tensor):
+                    frames_for_embed = att_mask.sum(dim=-1).cpu().tolist()
+                else:
+                    frames_for_embed = np.array(att_mask).sum(axis=-1).tolist()
+                audio_token_counts = [self.compute_audio_token_count(mel_frame_count) for mel_frame_count in
+                                      frames_for_embed]
+            if text is None:
+                text = [""] * num_samples
+            elif isinstance(text, str):
+                text = [text]
+            prompts = []
+            for idx in range(num_samples):
+                prompt = text[idx]
+                has_image = images is not None
+                audio_count = audio_token_counts[idx] if audio_token_counts is not None else None
+                prompt_str = prompt
+                if has_image:
+                    prompt_str = prompt_str.replace(self.boi_token, self.full_image_sequence)
+                if audio_count is not None:
+                    prompt_str = prompt_str.replace(self.boa_token, self.boa_token + (self.audio_token * audio_count))
+                prompts.append(prompt_str)
+            text_features = self.tokenizer(prompts, return_tensors=return_tensors, padding=True, truncation=True,
+                                           max_length=DEFAULT_MAX_LENGTH)
+            inputs = {**text_features}
+            if image_features:
+                inputs.update(image_features)
+            if audio_features:
+                inputs.update(audio_features)
+        if device is not None:
+            inputs = {k: v.to(device) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        return inputs
+
+    @property
+    def model_input_names(self) -> List[str]:
+        input_names = set()
+        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+            tokenizer_inputs = self.tokenizer.model_input_names
+            if isinstance(tokenizer_inputs, (list, set)):
+                input_names.update(tokenizer_inputs)
+            else:
+                input_names.add(str(tokenizer_inputs))
+            input_names.add("token_type_ids")
+        if hasattr(self, 'image_processor') and self.image_processor is not None:
+            image_inputs = self.image_processor.model_input_names
+            if isinstance(image_inputs, (list, set)):
+                input_names.update(image_inputs)
+            else:
+                input_names.add(str(image_inputs))
+        if hasattr(self, 'audio_processor') and self.audio_processor is not None:
+            audio_inputs = self.audio_processor.model_input_names
+            if isinstance(audio_inputs, (list, set)):
+                input_names.update(audio_inputs)
+            else:
+                input_names.add(str(audio_inputs))
+        return list(input_names)

special_tokens_map.json

@@ -0,0 +1,36 @@
+{
+  "boi_token": "<start_of_image>",
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<end_of_image>",
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<image_soft_token>",
+  "boa_token": "<start_of_audio>",
+  "eoa_token": "<end_of_audio>",
+  "audio_token": "<audio_soft_token>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88787c7cf85c7d14c8dd2a29cc86f69a1a7d151f306ce00bb54fe7dc35284b0e
+size 33384534

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
+size 4689074