| |
| |
| |
| |
| |
|
|
| from typing import Union, List |
|
|
| import torch |
|
|
| from torch import nn |
| from transformers.modeling_utils import PreTrainedModel |
|
|
| from .configuration_internvlu_generation_decoder import InternVLUGenerationDecoderConfig |
| from .internvlu_transformer import InternVLUTransformer2DModel |
|
|
|
|
| class InternVLUGenerationDecoder(PreTrainedModel): |
| """Transformer-based generation decoder used by the InternVLU diffusion stack.""" |
|
|
| config_class = InternVLUGenerationDecoderConfig |
| base_model_prefix = "" |
|
|
| def __init__(self, config: InternVLUGenerationDecoderConfig): |
| super().__init__(config) |
|
|
| self.config = config |
| config.decoder_config["video_position_scale_factor"] = ( |
| config.video_position_scale_factor |
| ) |
| config.decoder_config["vlm_cond_position_scale_factor"] = ( |
| config.vlm_cond_position_scale_factor |
| ) |
| self.decoder = InternVLUTransformer2DModel.from_config(config.decoder_config) |
|
|
| if self.config.padding_encoder_hidden_states: |
| self.encoder_padding_token = nn.Parameter( |
| torch.zeros(self.config.input_hidden_size) |
| ) |
|
|
| self.do_denormalize = True |
|
|
| self._init_decoder_projector() |
| self._init_visual_embeds_projector() |
|
|
| def _init_decoder_projector(self): |
| """Initialize the decoder_projector for multi model input.""" |
| if self.config.decoder_projector_type == "mlp2x_gelu": |
| self.decoder_projector = nn.Sequential( |
| nn.Linear( |
| self.config.input_hidden_size, |
| self.config.output_hidden_size * 3, |
| ), |
| nn.SiLU(), |
| nn.Linear( |
| self.config.output_hidden_size * 3, self.config.output_hidden_size |
| ), |
| ) |
| elif self.config.decoder_projector_type == "identity": |
| self.decoder_projector = nn.Identity() |
| else: |
| raise ValueError( |
| f"Unknown decoder_projector_type: {self.config.decoder_projector_type}" |
| ) |
|
|
| def _init_visual_embeds_projector(self): |
| """Initialize the visual_embeds_projector for visual encoder input.""" |
| self.visual_embeds_projector = nn.Identity() |
| return |
|
|
| @property |
| def device(self): |
| return self.decoder.device |
|
|
| def prepare_forward_input( |
| self, |
| encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]], |
| encoder_image_token_mask: Union[torch.Tensor, List[torch.Tensor]] = None, |
| **kwargs, |
| ): |
| return self._prepare_forward_input_default( |
| encoder_hidden_states, encoder_image_token_mask, **kwargs |
| ) |
|
|
| def _prepare_forward_input_default( |
| self, |
| encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]], |
| encoder_image_token_mask: Union[torch.Tensor, List[torch.Tensor]] = None, |
| **kwargs, |
| ): |
| """Project encoder hidden states and build attention masks (no visual embeds).""" |
| assert encoder_hidden_states is not None |
|
|
| if isinstance(encoder_hidden_states, list): |
| assert ( |
| self.config.padding_encoder_hidden_states |
| ), "encoder_hidden_states should be a list only when padding is enabled." |
| |
| |
| import os |
|
|
| pad_to_fix_length = self.config.pad_to_fix_length |
|
|
| if pad_to_fix_length: |
| max_length = self.config.max_sequence_length |
| else: |
| max_length = max([x.shape[0] for x in encoder_hidden_states]) |
| max_length = max( |
| max_length, self.config.max_sequence_length |
| ) |
|
|
| attention_masks = torch.zeros( |
| len(encoder_hidden_states), max_length, device=self.device |
| ).bool() |
| encoder_hidden_states_padded = self.encoder_padding_token[ |
| None, None |
| ].repeat(len(encoder_hidden_states), max_length, 1) |
| if encoder_image_token_mask is not None: |
| encoder_image_token_mask_padded = torch.zeros( |
| len(encoder_hidden_states), max_length, device=self.device |
| ).bool() |
| for idx, encoder_hidden_state in enumerate(encoder_hidden_states): |
| encoder_seq_len = min(encoder_hidden_state.shape[0], max_length) |
| encoder_hidden_states_padded[idx, :encoder_seq_len] = ( |
| encoder_hidden_state[:encoder_seq_len] |
| ) |
| attention_masks[idx, :encoder_seq_len] = True |
| if encoder_image_token_mask is not None: |
| encoder_image_token_mask_padded[idx, :encoder_seq_len] = ( |
| encoder_image_token_mask[idx][:encoder_seq_len] |
| ) |
| if encoder_image_token_mask is not None: |
| encoder_image_token_mask = encoder_image_token_mask_padded |
|
|
| encoder_hidden_states = encoder_hidden_states_padded |
|
|
| prefix_prompt_embeds = kwargs.pop("prefix_prompt_embeds", None) |
| if prefix_prompt_embeds is not None: |
| encoder_hidden_states = torch.concat( |
| [encoder_hidden_states, prefix_prompt_embeds], dim=1 |
| ) |
| |
| extended_attention_masks = torch.ones( |
| prefix_prompt_embeds.shape[0], |
| prefix_prompt_embeds.shape[1], |
| device=self.device, |
| ).bool() |
| attention_masks = torch.concat( |
| [attention_masks, extended_attention_masks], dim=1 |
| ) |
|
|
| encoder_hidden_states = self.decoder_projector(encoder_hidden_states) |
| return encoder_hidden_states, attention_masks, encoder_image_token_mask |
|
|
| def get_sigmas(self, timesteps, n_dim=4): |
| """Convert scheduler timesteps to sigma values with broadcast to `n_dim`.""" |
| sigmas = self.train_scheduler.sigmas.to(device=self.device, dtype=self.dtype) |
| schedule_timesteps = self.train_scheduler.timesteps.to(self.device) |
| timesteps = timesteps.to(self.device) |
| step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] |
|
|
| sigma = sigmas[step_indices].flatten() |
| while len(sigma.shape) < n_dim: |
| sigma = sigma.unsqueeze(-1) |
| return sigma |
|
|
