Upload Gemma3OmniForConditionalGeneration

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,132 @@
+{
+  "architectures": [
+    "Gemma3OmniForConditionalGeneration"
+  ],
+  "audio_token_index": 262151,
+  "auto_map": {
+    "AutoConfig": "configuration_gemma3_omni.Gemma3OmniConfig",
+    "AutoFeatureExtractor": "processing_gemma3_omni.Gemma3AudioFeatureExtractor",
+    "AutoModel": "modeling_gemma3_omni.Gemma3OmniForConditionalGeneration",
+    "AutoModelForCausalLM": "modeling_gemma3_omni.Gemma3OmniForConditionalGeneration",
+    "AutoProcessor": "processing_gemma3_omni.Gemma3OmniProcessor"
+  },
+  "boi_token_index": 255999,
+  "eoi_token_index": 256000,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_index": 262152,
+  "initializer_range": 0.02,
+  "mm_tokens_per_image": 256,
+  "model_type": "gemma3omni",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_logit_softcapping": null,
+    "final_logit_softcapping": null,
+    "head_dim": 128,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 5376,
+    "initializer_range": 0.02,
+    "intermediate_size": 21504,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention"
+    ],
+    "max_position_embeddings": 131072,
+    "model_type": "gemma3_text",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 62,
+    "num_key_value_heads": 16,
+    "query_pre_attn_scalar": 168,
+    "rms_norm_eps": 1e-06,
+    "rope_local_base_freq": 10000.0,
+    "rope_scaling": {
+      "factor": 8.0,
+      "rope_type": "linear"
+    },
+    "rope_theta": 1000000.0,
+    "sliding_window": 1024,
+    "torch_dtype": "bfloat16",
+    "use_cache": true,
+    "vocab_size": 262208
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.0",
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "image_size": 896,
+    "intermediate_size": 4304,
+    "layer_norm_eps": 1e-06,
+    "model_type": "siglip_vision_model",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "patch_size": 14,
+    "torch_dtype": "bfloat16",
+    "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_152,
+            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)

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "pad_token_id": 0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "4.53.0"
+}

model-00001-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:265edf8d47207aa1371ad2e9b48a198de6374fa649c185b902879c2c42e86303
+size 4922387560

model-00002-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ebcc142679e7a1982b4fbbbb1ef9a242b70eeb3731a3e98db15eadfe70296ba
+size 4954792944

model-00003-of-00012.safetensors

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

model-00004-of-00012.safetensors

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

model-00005-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11d1cdb15b770bf3a8c340f4124ed2def7262e4b77bc97be6b24ac4eabf377b9
+size 4954793016

model-00006-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6daf7a5de65a8ec47420d11fc4f66fe027873507009d2dedc490a9f9a441a53c
+size 4954793016

model-00007-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d18a18b1285e9d1d4f40d9a2df892b73c433116b0a5ac12abc59cf38bf345cd
+size 4954793016

model-00008-of-00012.safetensors

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

model-00009-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95f7b177c303ac266c1ff7fb897426235ac87ff275185971cee1462851b961ec
+size 4954793016

model-00010-of-00012.safetensors

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

model-00011-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8638f5b356e7b7f2d059365d9afd8dc03bdd2b524dee1c14338d62fa6eac483d
+size 4954793016

model-00012-of-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c042662255b2c23d6546318982a10474af91b1e2d0fdf83dcb265ed323ca4e5
+size 1288275520

modeling_gemma3_omni.py

@@ -0,0 +1,484 @@
+# -*- 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_list[-1])
+        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:
+        return token_type_ids[batch_idx, kv_idx] != 0
+    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.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    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).clone()
+
+        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.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    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:
+            if LigerFusedLinearCrossEntropyLoss is not None:
+                shift_hidden_states = hidden_states[..., :-1, :]  # (B, S-1, H)
+                shift_labels = labels[..., 1:]  # (B, S-1)
+                hidden_device = shift_hidden_states.device
+
+                if attention_mask is not None:
+                    shift_attention_mask = attention_mask[:, -shift_hidden_states.shape[1]:].to(hidden_device)
+                    shift_hidden_states = shift_hidden_states[shift_attention_mask != 0].contiguous()
+                    shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+                else:
+                    shift_hidden_states = shift_hidden_states.contiguous()
+                    shift_labels = shift_labels.contiguous()
+
+                shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size)  # (N, H)
+                shift_labels = shift_labels.view(-1).to(hidden_device)
+
+                loss_fct = LigerFusedLinearCrossEntropyLoss()
+                loss = loss_fct(self.lm_head.weight, shift_hidden_states, shift_labels)
+            else:
+                logits = logits.float()
+                shift_logits = logits[..., :-1, :]  # (B, S-1, V)
+                shift_labels = labels[..., 1:]  # (B, S-1)
+
+                if attention_mask is not None:
+                    shift_attention_mask = attention_mask[:, -shift_logits.shape[1]:].to(logits.device)
+                    shift_logits = shift_logits[shift_attention_mask != 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()
+
+                flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+                flat_labels = shift_labels.view(-1).to(shift_logits.device)
+
+                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",
+]