# 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")