| |
| """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 |
| |
| 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, |
| ): |
| |
| 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") |
|
|
