modeling_gemma3_omni.py

# -*- coding: utf-8 -*-
from __future__ import annotations

from typing import List, Optional, Tuple, Union

import torch
import torchaudio
from torch import nn
from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    Cache,
    Gemma3Config,
    PreTrainedModel,
    PretrainedConfig, StaticCache, HybridCache,
)
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.models.gemma3.modeling_gemma3 import (
    Gemma3CausalLMOutputWithPast,
    Gemma3ForConditionalGeneration,
    Gemma3RMSNorm,
)
from transformers.utils import is_torchdynamo_compiling, logging

from .speech_conformer_encoder import ConformerEncoder

logger = logging.get_logger(__name__)


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,
            audio_token_id: int = 0,
            **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.audio_token_id = audio_token_id


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 = config.audio_processor.get("config", None)
        encoder_config = {
            "activation": "swish",
            "activation_checkpointing": {
                "interval": 1,
                "module": "transformer",
                "offload": False
            },
            "attention_dim": 1024,
            "attention_heads": 16,
            "batch_norm": False,
            "bias_in_glu": True,
            "causal": True,
            "chunk_size": -1,
            "cnn_layer_norm": True,
            "conv_activation": "swish",
            "conv_glu_type": "swish",
            "depthwise_multiplier": 1,
            "depthwise_seperable_out_channel": 1024,
            "dropout_rate": 0.0,
            "encoder_embedding_config": {
                "input_size": 80
            },
            "ext_pw_kernel_size": 1,
            "ext_pw_out_channel": 1024,
            "input_layer": "nemo_conv",
            "input_size": 80,
            "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.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=config.sample_rate, n_mels=config.n_mels
        )
        self.proj = nn.Linear(self.encoder.ATTN_DIM, config.hidden_size, bias=False)
        self.layer_norm = nn.LayerNorm(config.hidden_size)
        self.post_init()

    # ---------- helpers ----------
    def wav2mel(self, wav: torch.Tensor) -> torch.Tensor:
        return self.mel(wav).clamp(min=1e-5).log().transpose(1, 2)

    # ---------- forward ----------
    @torch.no_grad()
    def forward(self, wav: torch.Tensor) -> torch.Tensor:  # (B,T) or (B,1,T)
        if wav.dim() == 3:
            wav = wav.squeeze(1)
        mel = self.wav2mel(wav)
        lengths = torch.full(
            (mel.size(0),), mel.size(1), dtype=torch.long, device=mel.device
        )
        hidden = self.encoder(mel, lengths)
        hidden = self.proj(hidden)
        return self.layer_norm(hidden)


# ──────────────────────────────────────────────────────────────────────────────
# Vision projector (與原版一致，只改 dtype)
# ──────────────────────────────────────────────────────────────────────────────
class Gemma3VisionProjector(nn.Module):
    def __init__(self, config: Gemma3Config):
        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)


# ──────────────────────────────────────────────────────────────────────────────
# Gemma-3 Multimodal wrapper
# ──────────────────────────────────────────────────────────────────────────────
class Gemma3OmniForConditionalGeneration(Gemma3ForConditionalGeneration):
    """Gemma-3 Omni：vision + audio + text causal LM."""

    def __init__(self, config: Gemma3Config):
        super().__init__(config)

        # ---- sub-modules
        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 = AutoModelForCausalLM.from_config(config=config.text_config)
        if language_model._tied_weights_keys is not None:
            self._tied_weights_keys = [f"language_model.{k}" for k in language_model._tied_weights_keys]
        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()

    # ---------- helper ----------
    def get_audio_features(self, audio_values: torch.Tensor) -> torch.Tensor:
        return self.audio_projector(audio_values)

    def _update_causal_mask(
            self,
            attention_mask,
            token_type_ids,
            past_key_values,
            cache_position,
            input_tensor,
            is_training: bool = False,
    ):
        if self.config.text_config._attn_implementation == "flash_attention_2":
            return attention_mask

        if attention_mask is not None and attention_mask.dim() == 4:
            # In this case we assume that the mask comes already in inverted
            # form and requires no inversion or slicing.
            return attention_mask

        using_static_cache = isinstance(past_key_values, StaticCache)
        min_dtype = torch.finfo(self.dtype).min
        inputs_lead_dim, sequence_length = input_tensor.shape[:2]
        if using_static_cache:
            target_length = past_key_values.get_max_cache_shape()
        elif isinstance(past_key_values, HybridCache):
            target_length = past_key_values.get_max_cache_shape()
        else:
            target_length = (
                attention_mask.shape[-1]
                if isinstance(attention_mask, torch.Tensor)
                else cache_position[0] + sequence_length + 1
            )

        if attention_mask is not None and attention_mask.dim() == 4:
            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
            return attention_mask

        causal_mask = torch.full(
            (sequence_length, target_length), fill_value=min_dtype, dtype=self.dtype, device=cache_position.device
        )

        # Causal diagonal mask only if training, otherwise attend to the whole prefix. Training-specific attn for prefix is handled below
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)

        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(inputs_lead_dim, 1, -1, -1)

        # Apply bidirectional mask on images if token type ids are provided
        if token_type_ids is not None and sequence_length != 1:
            token_type_mask = token_type_ids.unsqueeze(1) == token_type_ids.unsqueeze(2)
            token_type_mask[token_type_ids == 0] = False  # if text token do not change anything
            token_type_mask = token_type_mask.unsqueeze(1).to(causal_mask.device, dtype=torch.bool)
            causal_mask = causal_mask.clone()
            causal_mask[:, :, :, :sequence_length] = causal_mask[:, :, :, :sequence_length].masked_fill(
                token_type_mask, 0.0
            )

        if attention_mask is not None:
            causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
            mask_length = attention_mask.shape[-1]

            # Then apply padding mask (will mask pad tokens)
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
                padding_mask, min_dtype
            )

        return causal_mask

    # ---------- forward ----------
    def forward(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            pixel_values: Optional[torch.FloatTensor] = None,
            audio_values: Optional[torch.FloatTensor] = 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,
            logits_to_keep: Union[int, torch.Tensor] = 0,
            **lm_kwargs,
    ) -> Union[Tuple, Gemma3CausalLMOutputWithPast]:

        # === input validation ===
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("Exactly one of input_ids or inputs_embeds must be provided")

        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
        )

        is_training = token_type_ids is not None and labels is not None

        # OOV image token → pad
        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)

        # cache_position
        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,
            )

        # === merge image ===
        if pixel_values is not None:
            image_feat = self.get_image_features(pixel_values)
            special_image_mask = (
                (
                        inputs_embeds
                        == self.get_input_embeddings()(
                    torch.tensor(self.config.image_token_id, device=inputs_embeds.device)
                )
                )
                if input_ids is None
                else (
                        input_ids == self.config.image_token_id
                ).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            )
            if (
                    not is_torchdynamo_compiling()
                    and inputs_embeds[special_image_mask].numel() != image_feat.numel()
            ):
                raise ValueError("#image tokens ≠ #embedding slots")
            inputs_embeds = inputs_embeds.masked_scatter(
                special_image_mask, image_feat.to(inputs_embeds)
            )

        # === merge audio ===
        if audio_values is not None:
            audio_feat = self.get_audio_features(audio_values)
            # special_audio_mask = (
            #     (
            #         inputs_embeds
            #         == self.get_input_embeddings()(
            #             torch.tensor(self.config.audio_token_id, device=inputs_embeds.device)
            #         )
            #     )
            #     if input_ids is None
            #     else (
            #         input_ids == self.config.audio_token_id
            #     ).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            # )
            # if (
            #     not is_torchdynamo_compiling()
            #     and inputs_embeds[special_audio_mask].numel() != audio_feat.numel()
            # ):
            #     raise ValueError("#audio tokens ≠ #embedding slots")
            # inputs_embeds = inputs_embeds.masked_scatter(
            #     special_audio_mask, audio_feat.to(inputs_embeds)
            # )
            print(audio_feat.shape, inputs_embeds.shape)
            inputs_embeds = torch.cat([audio_feat, inputs_embeds], dim=1)

        # === label masking ===
        if labels is not None and self.pad_token_id in labels:
            logger.warning_once(
                "`labels` contains `pad_token_id`; they will be masked out at loss computation."
            )
            labels = torch.where(
                input_ids == self.pad_token_id, self.config.ignore_index, labels
            )

        causal_mask = self._update_causal_mask(
            attention_mask,
            token_type_ids,
            past_key_values,
            cache_position,
            inputs_embeds,
            is_training,
        )

        outputs: CausalLMOutputWithPast = self.language_model(
            attention_mask=causal_mask,
            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,
            cache_position=cache_position,
            logits_to_keep=logits_to_keep,
            **lm_kwargs,
        )

        # === loss ===
        logits = outputs.logits
        loss = None
        if labels is not None:
            logits = logits.float()
            shift_logits = logits[..., :-1, :]
            shift_labels = labels[..., 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 != 0].contiguous()
            loss = nn.CrossEntropyLoss()(
                shift_logits.view(-1, self.config.text_config.vocab_size),
                shift_labels.view(-1),
            )

        return Gemma3CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            image_hidden_states=image_feat if pixel_values is not None else None,
        )


# ──────────────────────────────────────────────────────────────────────────────
# exports
# ──────────────────────────────────────────────────────────────────────────────
__all__ = [
    "Gemma3OmniForConditionalGeneration",
]