modeling_premonition.py

# Copyright 2026 Prem Research. Licensed under the Apache License, Version 2.0.
"""Gemma 4 audio/video prefill for the unchanged DiffusionGemma decoder and sampler."""

import torch
from torch import nn
from transformers import AutoModel
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.diffusion_gemma.modeling_diffusion_gemma import (
    DiffusionGemmaDecoderModel,
    DiffusionGemmaEncoderModel,
    DiffusionGemmaForBlockDiffusion,
    DiffusionGemmaModel,
    DiffusionGemmaPreTrainedModel,
)
from transformers.models.gemma4.modeling_gemma4 import Gemma4Model, Gemma4MultimodalEmbedder

from .configuration_premonition import PremonitionConfig


class PremonitionEncoderModel(DiffusionGemmaEncoderModel):
    config_class = PremonitionConfig
    input_modalities = ("text", "image", "video", "audio")

    def __init__(self, config):
        super().__init__(config)
        self.audio_tower = (
            AutoModel.from_config(config.audio_config) if config.audio_config is not None else None
        )
        self.embed_audio = (
            Gemma4MultimodalEmbedder(config.audio_config, config.text_config)
            if config.audio_config is not None
            else None
        )
        self.post_init()

    get_audio_features = Gemma4Model.get_audio_features
    get_video_features = Gemma4Model.get_video_features

    def forward(
        self,
        input_ids=None,
        pixel_values=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        mm_token_type_ids=None,
        inputs_embeds=None,
        image_position_ids=None,
        input_features=None,
        input_features_mask=None,
        pixel_values_videos=None,
        video_position_ids=None,
        **kwargs,
    ):
        if (input_features is None) != (input_features_mask is None):
            raise ValueError("Audio requires both input_features and input_features_mask.")
        if input_features is None and pixel_values_videos is None:
            return super().forward(
                input_ids=input_ids,
                pixel_values=pixel_values,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                mm_token_type_ids=mm_token_type_ids,
                inputs_embeds=inputs_embeds,
                image_position_ids=image_position_ids,
                **kwargs,
            )
        if input_ids is None or inputs_embeds is not None:
            raise ValueError("Audio/video prefill requires input_ids, without inputs_embeds.")
        if pixel_values_videos is not None and video_position_ids is None:
            raise ValueError("Video requires video_position_ids.")

        masks = {
            "image": input_ids == self.config.image_token_id,
            "video": input_ids == self.config.video_token_id,
            "audio": input_ids == self.config.audio_token_id,
        }
        multimodal_mask = masks["image"] | masks["video"] | masks["audio"]
        safe_ids = input_ids.masked_fill(multimodal_mask, self.config.text_config.pad_token_id)
        inputs_embeds = self.get_input_embeddings()(safe_ids)

        def merge(modality, features):
            mask = masks[modality].unsqueeze(-1).expand_as(inputs_embeds)
            if int(masks[modality].sum()) != features.numel() // inputs_embeds.shape[-1]:
                raise ValueError(
                    f"{modality} placeholder count ({int(masks[modality].sum())}) "
                    f"does not match valid feature count ({features.shape[0]})."
                )
            return inputs_embeds.masked_scatter(mask, features.to(inputs_embeds))

        if pixel_values is not None:
            features = self.get_image_features(
                pixel_values, image_position_ids, return_dict=True
            ).pooler_output
            inputs_embeds = merge("image", features)
        if pixel_values_videos is not None:
            features = self.get_video_features(
                pixel_values_videos, video_position_ids, return_dict=True
            ).pooler_output
            inputs_embeds = merge("video", torch.cat(features, dim=0))
        if input_features is not None:
            audio_outputs = self.get_audio_features(
                input_features, input_features_mask, return_dict=True
            )
            features = audio_outputs.pooler_output
            features = features[audio_outputs.attention_mask.to(features.device)]
            inputs_embeds = merge("audio", features)

        if position_ids is None:
            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = (
                torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen
            ).unsqueeze(0)
        if not isinstance(attention_mask, dict):
            attention_mask = self.create_masks_for_generate(
                config=self.config.get_text_config(),
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                position_ids=position_ids,
                mm_token_type_ids=mm_token_type_ids,
            )
        return_dict = kwargs.pop("return_dict", True)
        outputs = self.language_model(
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            return_dict=True,
            **kwargs,
        )
        result = BaseModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
        return result if return_dict else result.to_tuple()


class PremonitionModel(DiffusionGemmaModel):
    config_class = PremonitionConfig
    input_modalities = ("text", "image", "video", "audio")

    def __init__(self, config):
        DiffusionGemmaPreTrainedModel.__init__(self, config)
        self.encoder = PremonitionEncoderModel(config)
        self.decoder = DiffusionGemmaDecoderModel(config)
        self.post_init()


class PremonitionForBlockDiffusion(DiffusionGemmaForBlockDiffusion):
    config_class = PremonitionConfig
    input_modalities = ("text", "image", "video", "audio")

    def __init__(self, config):
        DiffusionGemmaPreTrainedModel.__init__(self, config)
        self.model = PremonitionModel(config)
        self.lm_head = nn.Linear(
            config.text_config.hidden_size, config.text_config.vocab_size, bias=False
        )
        self.final_logit_softcapping = config.text_config.final_logit_softcapping
        self.post_init()
        if config.freeze_backbone:
            self.freeze_for_projector_training()

    def freeze_for_projector_training(self):
        self.requires_grad_(False)
        projector = self.model.encoder.embed_audio
        if projector is not None:
            projector.requires_grad_(True)
        return [name for name, param in self.named_parameters() if param.requires_grad]

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        loaded = super().from_pretrained(*args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded
        # HF's meta-device loader replaces Parameters and can reset requires_grad.
        if model.config.freeze_backbone:
            model.freeze_for_projector_training()
        return loaded

    def _prepare_denoiser_inputs(
        self,
        decoder_attention_mask,
        past_key_values,
        sampler,
        diffusion_stopping_criteria,
        batch_size,
        device,
        model_kwargs,
    ):
        # Consume modalities once, during prefill. Never re-encode on subsequent canvases.
        for key in (
            "input_features",
            "input_features_mask",
            "pixel_values_videos",
            "video_position_ids",
        ):
            model_kwargs.pop(key, None)
        return super()._prepare_denoiser_inputs(
            decoder_attention_mask,
            past_key_values,
            sampler,
            diffusion_stopping_criteria,
            batch_size,
            device,
            model_kwargs,
        )


PremonitionModel.register_for_auto_class("AutoModel")
PremonitionForBlockDiffusion.register_for_auto_class("AutoModelForImageTextToText")