|
|
|
from typing import List, Optional, Tuple, Union |
|
import torch, math |
|
import torch.utils.checkpoint |
|
from torch import nn |
|
import transformers |
|
from flash_attn import flash_attn_varlen_func |
|
from transformers.activations import ACT2FN |
|
from PIL import Image |
|
import io, fire |
|
from torch.nn import functional as F |
|
|
|
class OmniVisualEncoder(transformers.models.qwen2_vl.modeling_qwen2_vl.Qwen2VisionTransformerPretrainedModel): |
|
def __init__(self, config): |
|
super().__init__(config) |
|
self.config_attn_implementation = 'flash_attention_2' |
|
self.gradient_checkpointing = True |
|
self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint |
|
self.merge_size = config.merge_size if hasattr(config, 'merge_size') else 2 |
|
del self.merger |
|
|
|
def forward( |
|
self, |
|
pixel_values: torch.Tensor, |
|
grid_thw: torch.Tensor, |
|
): |
|
hidden_states = pixel_values.to(self.get_dtype()) |
|
grid_thw = grid_thw.to(pixel_values.device) |
|
|
|
hidden_states = self.patch_embed(hidden_states) |
|
rotary_pos_emb = self.rot_pos_emb(grid_thw) |
|
|
|
cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( |
|
dim=0, dtype=torch.int32 |
|
) |
|
cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) |
|
|
|
for blk in self.blocks: |
|
if self.gradient_checkpointing and self.training: |
|
hidden_states = self._gradient_checkpointing_func(blk.__call__, hidden_states, cu_seqlens, rotary_pos_emb) |
|
else: |
|
hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb) |
|
|
|
return hidden_states |
|
|
|
@torch.no_grad() |
|
def fake_input(self, device): |
|
merge_size = max(self.merge_size, self.config.spatial_merge_size) |
|
fake_image = torch.zeros([ |
|
1, |
|
self.config.temporal_patch_size, |
|
3, |
|
merge_size // self.config.spatial_merge_size, |
|
self.config.spatial_merge_size, |
|
self.config.patch_size, |
|
merge_size // self.config.spatial_merge_size, |
|
self.config.spatial_merge_size, |
|
self.config.patch_size, |
|
], dtype=torch.float32, device=device) |
|
patches = fake_image.permute(0, 3, 6, 4, 7, 2, 1, 5, 8) |
|
flatten_patches = patches.reshape( |
|
merge_size * merge_size, 3 * self.config.temporal_patch_size * self.config.patch_size * self.config.patch_size |
|
) |
|
return [flatten_patches], [(1, merge_size, merge_size)], [1] |
|
|
|
|
|
class OmniVisualBridge(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
self.config = config |
|
self.merge_size = self.config.merge_size if hasattr(self.config, 'merge_size') else 2 |
|
self.hidden_size = config.embed_dim * (self.merge_size**2) |
|
self.ln_q = nn.LayerNorm(config.embed_dim, eps=1e-6) |
|
self.mlp = nn.Sequential( |
|
nn.Linear(self.hidden_size, self.hidden_size), |
|
nn.GELU(), |
|
nn.Linear(self.hidden_size, config.hidden_size), |
|
) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) |
|
return x |
|
|
|
|
|
if __name__ == '__main__': |
|
fire.Fire() |
|
|
