File size: 6,734 Bytes
a130d59 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | # --------------------------------------------------------
# InternVL-U
# Copyright (c) 2026 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
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."
# if padding is enabled, we need to pad the encoder_hidden_states
# if self.config.padding_encoder_hidden_states:
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
) # at least 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
)
# attention mask
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
|