File size: 7,499 Bytes
c6bc767 45cf443 c6bc767 45cf443 92c7321 45cf443 c6bc767 45cf443 c6bc767 45cf443 c6bc767 45cf443 c6bc767 45cf443 c6bc767 45cf443 | 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 167 | import os
import torch
import torch.nn.functional as F
import warnings
from typing import Optional, Tuple, List, Union
from torch import nn
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
from transformers import SiglipVisionModel, SiglipImageProcessor, logging as hf_logging
from core import precompute_freqs_cis, MOEFeedForward
from models.lm.model import LMForCausalLM
from models.vlm.config import VLMConfig
from encoders.vision import SiglipVisionEncoder
from projectors import MMVisionProjector
warnings.filterwarnings('ignore')
class VLM(LMForCausalLM):
config_class = VLMConfig
def __init__(self, config: VLMConfig = None, vision_model_path: Optional[str] = None):
self.config = config or VLMConfig()
super().__init__(self.config)
self.vision_encoder = SiglipVisionEncoder(vision_model_path) if vision_model_path else SiglipVisionEncoder()
self.vision_proj = MMVisionProjector(
self.config.image_hidden_size, self.config.hidden_size, target_tokens=self.config.image_token_len
)
@staticmethod
def get_vision_model(model_path: str):
hf_logging.set_verbosity_error()
if not os.path.exists(model_path):
return None, None
try:
model = SiglipVisionModel.from_pretrained(model_path)
except (RuntimeError, ValueError):
return None, None
processor = SiglipImageProcessor.from_pretrained(model_path)
for param in model.parameters():
param.requires_grad = False
return model.eval(), processor
@staticmethod
def image2tensor(image, processor):
if image.mode in ['RGBA', 'LA']:
image = image.convert('RGB')
return processor(images=image, return_tensors="pt")
@staticmethod
def get_image_embeddings(image_inputs, vision_model):
if vision_model is None:
return None
return vision_model.encode(image_inputs)
@torch.compiler.disable
def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
if vision_tensors is None or not self.config.image_ids:
return h
marker, vf = self.config.image_ids[0], vision_tensors
if vf.dim() == 3:
vf = vf.unsqueeze(1)
out = []
for b in range(h.size(0)):
hb, seq, k, i = h[b], tokens[b].tolist(), 0, 0
while i < len(seq):
if seq[i] == marker:
start = i
while i < len(seq) and seq[i] == marker:
i += 1
if k < vf.size(1):
hb = torch.cat((hb[:start], vf[b][k][:i - start], hb[i:]), dim=0)[:seqlen]
k += 1
else:
i += 1
out.append(hb)
return torch.stack(out)
def forward(self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
use_cache: bool = False,
logits_to_keep: Union[int, torch.Tensor] = 0,
labels: Optional[torch.Tensor] = None,
pixel_values: Optional[torch.FloatTensor] = None,
**args):
batch_size, seq_length = input_ids.shape
if hasattr(past_key_values, 'layers'):
past_key_values = None
past_key_values = past_key_values or [None] * len(self.model.layers)
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
if pixel_values is not None and start_pos == 0:
if hasattr(pixel_values, 'keys'):
sample_val = next(iter(pixel_values.values()))
if sample_val.ndim == 5:
bs, num = sample_val.shape[:2]
vision_tensors = self.vision_proj(
VLM.get_image_embeddings(
{k: v.flatten(0, 1) for k, v in pixel_values.items()}, self.vision_encoder
)
).view(bs, num, self.config.image_token_len, -1)
else:
vision_tensors = self.vision_proj(
VLM.get_image_embeddings(pixel_values, self.vision_encoder)
)
else:
if len(pixel_values.shape) == 6:
pixel_values = pixel_values.squeeze(2)
bs, num, c, im_h, im_w = pixel_values.shape
vision_tensors = torch.stack(
[self.vision_proj(VLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder))
for i in range(num)], dim=1
)
hidden_states = self.count_vision_proj(
tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors, seqlen=input_ids.shape[1]
)
if self.model.freqs_cos[0, 0] == 0:
freqs_cos, freqs_sin = precompute_freqs_cis(
dim=self.config.head_dim, end=self.config.max_position_embeddings,
rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling
)
self.model.freqs_cos, self.model.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
position_embeddings = (
self.model.freqs_cos[start_pos:start_pos + seq_length],
self.model.freqs_sin[start_pos:start_pos + seq_length]
)
presents = []
for layer_idx, (layer, past_key_value) in enumerate(zip(self.model.layers, past_key_values)):
hidden_states, present = layer(
hidden_states, position_embeddings,
past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask
)
presents.append(present)
hidden_states = self.model.norm(hidden_states)
aux_loss = sum([l.mlp.aux_loss for l in self.model.layers if isinstance(l.mlp, MOEFeedForward)],
hidden_states.new_zeros(1).squeeze())
aux_loss = aux_loss + sum(p.sum() for p in self.vision_proj.parameters()) * 0
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
loss = None
if labels is not None:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
output = MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=presents, hidden_states=hidden_states)
return output
def generate(self, *args, num_return_sequences=1, **kwargs):
if num_return_sequences > 1 and 'pixel_values' in kwargs:
pv = kwargs['pixel_values']
if hasattr(pv, 'keys'):
kwargs['pixel_values'] = {k: v.repeat(num_return_sequences, *([1] * (v.ndim - 1))) for k, v in pv.items()}
else:
kwargs['pixel_values'] = pv.repeat(num_return_sequences, *([1] * (pv.ndim - 1)))
return super().generate(*args, num_return_sequences=num_return_sequences, **kwargs)
|