model: support NVILA and NVILA Lite (#10399)
This commit is contained in:
Zijian Zhang
@@ -914,12 +914,13 @@ multimodal_model_archs = [
|
||||
"InternVLChatModel",
|
||||
"InternS1ForConditionalGeneration",
|
||||
"Phi4MMForCausalLM",
|
||||
"VILAForConditionalGeneration",
|
||||
"Step3VLForConditionalGeneration",
|
||||
"POINTSV15ChatModel",
|
||||
"DotsVLMForCausalLM",
|
||||
"DotsOCRForCausalLM",
|
||||
"Sarashina2VisionForCausalLM",
|
||||
"NVILAForConditionalGeneration",
|
||||
"NVILALiteForConditionalGeneration",
|
||||
"DeepseekOCRForCausalLM",
|
||||
]
|
||||
|
||||
|
||||
355
python/sglang/srt/models/nvila.py
Normal file
355
python/sglang/srt/models/nvila.py
Normal file
@@ -0,0 +1,355 @@
|
||||
import itertools
|
||||
import math
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import einops
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
||||
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||
from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
|
||||
|
||||
import sglang.srt.managers.mm_utils as mm_utils
|
||||
import sglang.srt.model_loader.weight_utils as weight_utils
|
||||
import sglang.srt.utils as utils
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.models.qwen2 import Qwen2ForCausalLM
|
||||
|
||||
MM_HIDDEN_SIZE = 3456
|
||||
|
||||
|
||||
class NVILAConfig(PretrainedConfig):
|
||||
model_type = "nvila"
|
||||
sub_configs = {
|
||||
"text_config": Qwen2Config,
|
||||
"vision_config": SiglipVisionConfig,
|
||||
}
|
||||
_auto_class = "AutoConfig"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
text_config: dict[str, Any] | None = None,
|
||||
vision_config: dict[str, Any] | None = None,
|
||||
image_token_id: int | None = None,
|
||||
video_token_id: int | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.text_config = (
|
||||
Qwen2Config(**text_config) if text_config is not None else Qwen2Config()
|
||||
)
|
||||
self.vision_config = (
|
||||
SiglipVisionConfig(**vision_config)
|
||||
if vision_config is not None
|
||||
else SiglipVisionConfig()
|
||||
)
|
||||
|
||||
self.image_token_id = image_token_id if image_token_id is not None else -1
|
||||
self.video_token_id = video_token_id if video_token_id is not None else -1
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
class NVILAMultiModalProjectorDownsampleBlock(nn.Module):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
batch_size, sequence_length, hidden_size = x.shape
|
||||
|
||||
feat_size = math.isqrt(sequence_length)
|
||||
|
||||
features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
|
||||
|
||||
pad_after = feat_size % 2
|
||||
if pad_after > 0:
|
||||
features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
|
||||
feat_size = feat_size + pad_after
|
||||
|
||||
features = features.reshape(
|
||||
batch_size, feat_size // 2, 2, feat_size // 2, 2, hidden_size
|
||||
)
|
||||
features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
|
||||
features = features.reshape(batch_size, -1, 4 * hidden_size)
|
||||
|
||||
return features
|
||||
|
||||
|
||||
class NVILAMultiModalProjector(nn.Module):
|
||||
def __init__(self, config: NVILAConfig):
|
||||
super().__init__()
|
||||
|
||||
self.layers = nn.Sequential(
|
||||
NVILAMultiModalProjectorDownsampleBlock(),
|
||||
nn.LayerNorm(MM_HIDDEN_SIZE * 4),
|
||||
nn.Linear(MM_HIDDEN_SIZE * 4, config.text_config.hidden_size),
|
||||
nn.GELU(),
|
||||
nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size),
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class NVILAForConditionalGeneration(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: NVILAConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
|
||||
self.vision_tower = SiglipVisionModel(config.vision_config)
|
||||
self.mm_projector = NVILAMultiModalProjector(config)
|
||||
self.llm = Qwen2ForCausalLM(
|
||||
config=config.text_config,
|
||||
quant_config=quant_config,
|
||||
prefix=utils.add_prefix("llm", prefix),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
positions: Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
get_embedding: bool = False,
|
||||
) -> LogitsProcessorOutput:
|
||||
output = mm_utils.general_mm_embed_routine(
|
||||
input_ids=input_ids,
|
||||
forward_batch=forward_batch,
|
||||
language_model=self.llm,
|
||||
data_embedding_funcs={
|
||||
Modality.IMAGE: self.get_image_feature,
|
||||
Modality.VIDEO: self.get_image_feature,
|
||||
},
|
||||
get_embedding=get_embedding,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
assert isinstance(output, LogitsProcessorOutput)
|
||||
|
||||
return output
|
||||
|
||||
def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor:
|
||||
block_sizes = (
|
||||
list(
|
||||
itertools.chain.from_iterable(
|
||||
x.block_sizes for x in mm_input if hasattr(x, "block_sizes")
|
||||
)
|
||||
)
|
||||
or None
|
||||
)
|
||||
pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0)
|
||||
|
||||
vision_tower_output: BaseModelOutputWithPooling = self.vision_tower(
|
||||
pixel_values.to(
|
||||
device=self.vision_tower.device, dtype=self.vision_tower.dtype
|
||||
),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
assert vision_tower_output.hidden_states is not None
|
||||
|
||||
vision_features: Tensor = vision_tower_output.hidden_states[-2]
|
||||
|
||||
vision_features_list, block_sizes = merge_features_for_dynamic_s2(
|
||||
vision_features,
|
||||
block_sizes=(
|
||||
block_sizes
|
||||
if block_sizes is not None
|
||||
else [None] * vision_features.shape[0]
|
||||
),
|
||||
resize_output_to_scale_idx=-1,
|
||||
scales=[448, 896, 1344],
|
||||
)
|
||||
|
||||
vision_features_list = [
|
||||
split_chessboard(x, block_size[0], block_size[1])
|
||||
for x, block_size in zip(vision_features_list, block_sizes)
|
||||
]
|
||||
|
||||
vision_features = torch.cat(
|
||||
[einops.rearrange(x, "b c h w -> b (h w) c") for x in vision_features_list]
|
||||
)
|
||||
|
||||
vision_features = self.mm_projector(vision_features)
|
||||
|
||||
vision_features_list = list(
|
||||
vision_features.split(
|
||||
[block_size[0] * block_size[1] for block_size in block_sizes], dim=0
|
||||
)
|
||||
)
|
||||
vision_features_list = [
|
||||
merge_chessboard(x, block_size[0], block_size[1])
|
||||
for x, block_size in zip(vision_features_list, block_sizes)
|
||||
]
|
||||
|
||||
vision_features = torch.stack(
|
||||
[einops.rearrange(x, "1 c h w -> (h w) c") for x in vision_features_list]
|
||||
)
|
||||
|
||||
vision_features = einops.rearrange(vision_features, "n p d -> (n p) d")
|
||||
|
||||
return vision_features
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None:
|
||||
params_dict = dict(self.named_parameters())
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if name.startswith("llm."):
|
||||
self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
|
||||
else:
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", weight_utils.default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
def pad_input_ids(
|
||||
self, input_ids: list[int], mm_inputs: MultimodalInputs
|
||||
) -> list[int]:
|
||||
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||
return pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||
|
||||
|
||||
def merge_chessboard(x, num_split_h, num_split_w):
|
||||
"""
|
||||
x: b * n * c or b * h * w * c
|
||||
out: b * c * h * w
|
||||
Assuming x contains num_split**2 sub-squares concatenated along batch dimension, merge the sub-squares back to the original whole square.
|
||||
"""
|
||||
B = x.shape[0]
|
||||
if x.dim() == 3:
|
||||
N = x.shape[1]
|
||||
x = einops.rearrange(
|
||||
x, "b (h w) c -> b c h w", h=math.isqrt(N), w=math.isqrt(N)
|
||||
)
|
||||
|
||||
assert B % (num_split_h * num_split_w) == 0
|
||||
b = B // (num_split_h * num_split_w)
|
||||
|
||||
x_merge = torch.cat(
|
||||
[
|
||||
torch.cat(
|
||||
[
|
||||
x[(i * num_split_w + j) * b : (i * num_split_w + j + 1) * b]
|
||||
for j in range(num_split_w)
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
for i in range(num_split_h)
|
||||
],
|
||||
dim=-2,
|
||||
)
|
||||
|
||||
return x_merge
|
||||
|
||||
|
||||
def merge_features_for_dynamic_s2(
|
||||
image_features, block_sizes, *, scales, resize_output_to_scale_idx
|
||||
):
|
||||
image_features_each_image = []
|
||||
new_block_sizes = []
|
||||
block_cnt = 0
|
||||
for block_size_each_image in block_sizes:
|
||||
if block_size_each_image is None:
|
||||
cur_features = image_features[block_cnt : block_cnt + 1]
|
||||
cur_features = einops.rearrange(
|
||||
cur_features,
|
||||
"1 (h w) c -> 1 c h w",
|
||||
h=math.isqrt(cur_features.shape[1]),
|
||||
)
|
||||
cur_features = cur_features.repeat(1, len(scales), 1, 1)
|
||||
image_features_each_image.append(cur_features)
|
||||
new_block_sizes.append((1, 1))
|
||||
block_cnt += 1
|
||||
else:
|
||||
cur_features_each_scale = []
|
||||
for scale in scales[:-1]:
|
||||
num_blocks_this_scale = (scale // scales[0]) ** 2
|
||||
cur_features_each_scale.append(
|
||||
merge_chessboard(
|
||||
image_features[block_cnt : block_cnt + num_blocks_this_scale],
|
||||
num_split_h=scale // scales[0],
|
||||
num_split_w=scale // scales[0],
|
||||
)
|
||||
) # 1 * C * H * W
|
||||
block_cnt += num_blocks_this_scale
|
||||
num_blocks_last_scale = block_size_each_image[0] * block_size_each_image[1]
|
||||
cur_features_each_scale.append(
|
||||
merge_chessboard(
|
||||
image_features[block_cnt : block_cnt + num_blocks_last_scale],
|
||||
num_split_h=block_size_each_image[0],
|
||||
num_split_w=block_size_each_image[1],
|
||||
)
|
||||
) # 1 * C * H * W
|
||||
block_cnt += num_blocks_last_scale
|
||||
|
||||
# resize and concat features from different scales
|
||||
output_size = cur_features_each_scale[resize_output_to_scale_idx].shape[-2:]
|
||||
cur_features = torch.cat(
|
||||
[
|
||||
F.interpolate(
|
||||
cur_features_each_scale[i].to(torch.float32),
|
||||
size=output_size,
|
||||
mode="area",
|
||||
).to(cur_features_each_scale[i].dtype)
|
||||
for i in range(len(cur_features_each_scale))
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
image_features_each_image.append(cur_features)
|
||||
|
||||
if (
|
||||
resize_output_to_scale_idx == len(scales) - 1
|
||||
or resize_output_to_scale_idx == -1
|
||||
):
|
||||
new_block_sizes.append(block_size_each_image)
|
||||
else:
|
||||
new_block_sizes.append(
|
||||
(
|
||||
scales[resize_output_to_scale_idx] // scales[0],
|
||||
scales[resize_output_to_scale_idx] // scales[0],
|
||||
)
|
||||
)
|
||||
|
||||
assert block_cnt == len(
|
||||
image_features
|
||||
), f"The number of blocks ({block_cnt}) does not match length of image_features ({len(image_features)})!"
|
||||
|
||||
return image_features_each_image, new_block_sizes
|
||||
|
||||
|
||||
def split_chessboard(x, num_split_h, num_split_w):
|
||||
"""
|
||||
x: b * c * h * w
|
||||
out: b * c * h * w
|
||||
Deividing x into num_split**2 sub-squares, and concatenate all the sub-squares on the batch dimension
|
||||
"""
|
||||
B, C, H, W = x.shape
|
||||
assert H % num_split_h == 0 and W % num_split_w == 0
|
||||
h, w = H // num_split_h, W // num_split_w
|
||||
x_split = torch.cat(
|
||||
[
|
||||
x[:, :, i * h : (i + 1) * h, j * w : (j + 1) * w]
|
||||
for i in range(num_split_h)
|
||||
for j in range(num_split_w)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
return x_split
|
||||
|
||||
|
||||
EntryClass = [NVILAForConditionalGeneration]
|
||||
184
python/sglang/srt/models/nvila_lite.py
Normal file
184
python/sglang/srt/models/nvila_lite.py
Normal file
@@ -0,0 +1,184 @@
|
||||
import math
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import einops
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
||||
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||
from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
|
||||
|
||||
import sglang.srt.managers.mm_utils as mm_utils
|
||||
import sglang.srt.model_loader.weight_utils as weight_utils
|
||||
import sglang.srt.utils as utils
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.models.qwen2 import Qwen2ForCausalLM
|
||||
|
||||
MM_HIDDEN_SIZE = 1152
|
||||
|
||||
|
||||
class NVILALiteConfig(PretrainedConfig):
|
||||
model_type = "nvila_lite"
|
||||
sub_configs = {
|
||||
"text_config": Qwen2Config,
|
||||
"vision_config": SiglipVisionConfig,
|
||||
}
|
||||
_auto_class = "AutoConfig"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
text_config: dict[str, Any] | None = None,
|
||||
vision_config: dict[str, Any] | None = None,
|
||||
image_token_id: int | None = None,
|
||||
video_token_id: int | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.text_config = (
|
||||
Qwen2Config(**text_config) if text_config is not None else Qwen2Config()
|
||||
)
|
||||
self.vision_config = (
|
||||
SiglipVisionConfig(**vision_config)
|
||||
if vision_config is not None
|
||||
else SiglipVisionConfig()
|
||||
)
|
||||
|
||||
self.image_token_id = image_token_id if image_token_id is not None else -1
|
||||
self.video_token_id = video_token_id if video_token_id is not None else -1
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
class NVILALiteMultiModalProjectorDownsampleBlock(nn.Module):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
batch_size, sequence_length, hidden_size = x.shape
|
||||
|
||||
feat_size = math.isqrt(sequence_length)
|
||||
|
||||
features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
|
||||
|
||||
pad_after = (3 - feat_size % 3) % 3
|
||||
if pad_after > 0:
|
||||
features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
|
||||
feat_size = feat_size + pad_after
|
||||
|
||||
features = features.reshape(
|
||||
batch_size, feat_size // 3, 3, feat_size // 3, 3, hidden_size
|
||||
)
|
||||
features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
|
||||
features = features.reshape(batch_size, -1, 9 * hidden_size)
|
||||
|
||||
return features
|
||||
|
||||
|
||||
class NVILALiteMultiModalProjector(nn.Module):
|
||||
def __init__(self, config: NVILALiteConfig):
|
||||
super().__init__()
|
||||
|
||||
self.layers = nn.Sequential(
|
||||
NVILALiteMultiModalProjectorDownsampleBlock(),
|
||||
nn.LayerNorm(MM_HIDDEN_SIZE * 9),
|
||||
nn.Linear(MM_HIDDEN_SIZE * 9, MM_HIDDEN_SIZE * 3),
|
||||
nn.GELU(),
|
||||
nn.LayerNorm(MM_HIDDEN_SIZE * 3),
|
||||
nn.Linear(MM_HIDDEN_SIZE * 3, config.text_config.hidden_size),
|
||||
nn.GELU(),
|
||||
nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size),
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class NVILALiteForConditionalGeneration(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: NVILALiteConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
|
||||
self.vision_tower = SiglipVisionModel(config.vision_config)
|
||||
self.mm_projector = NVILALiteMultiModalProjector(config)
|
||||
self.llm = Qwen2ForCausalLM(
|
||||
config=config.text_config,
|
||||
quant_config=quant_config,
|
||||
prefix=utils.add_prefix("llm", prefix),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
positions: Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
get_embedding: bool = False,
|
||||
) -> LogitsProcessorOutput:
|
||||
output = mm_utils.general_mm_embed_routine(
|
||||
input_ids=input_ids,
|
||||
forward_batch=forward_batch,
|
||||
language_model=self.llm,
|
||||
data_embedding_funcs={
|
||||
Modality.IMAGE: self.get_image_feature,
|
||||
Modality.VIDEO: self.get_image_feature,
|
||||
},
|
||||
get_embedding=get_embedding,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
assert isinstance(output, LogitsProcessorOutput)
|
||||
|
||||
return output
|
||||
|
||||
def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor:
|
||||
pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0)
|
||||
|
||||
vision_tower_output: BaseModelOutputWithPooling = self.vision_tower(
|
||||
pixel_values,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
assert vision_tower_output.hidden_states is not None
|
||||
|
||||
vision_features = vision_tower_output.hidden_states[-2]
|
||||
|
||||
vision_features = self.mm_projector(vision_features)
|
||||
|
||||
vision_features = einops.rearrange(vision_features, "n p d -> (n p) d")
|
||||
|
||||
return vision_features
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None:
|
||||
params_dict = dict(self.named_parameters())
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if name.startswith("llm."):
|
||||
self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
|
||||
else:
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", weight_utils.default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
def pad_input_ids(
|
||||
self, input_ids: list[int], mm_inputs: MultimodalInputs
|
||||
) -> list[int]:
|
||||
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||
return pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||
|
||||
|
||||
EntryClass = [NVILALiteForConditionalGeneration]
|
||||
@@ -1,306 +0,0 @@
|
||||
import logging
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, cast
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
||||
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||
from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
|
||||
|
||||
import sglang.srt.managers.mm_utils as mm_utils
|
||||
import sglang.srt.model_loader.weight_utils as weight_utils
|
||||
import sglang.srt.utils as utils
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
|
||||
from sglang.srt.layers.pooler import Pooler, PoolingType
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.models.qwen2 import Qwen2ForCausalLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
##### BEGIN COPY configuration.py #####
|
||||
|
||||
|
||||
class VILAConfig(PretrainedConfig):
|
||||
# Class attributes.
|
||||
model_type: str = "vila"
|
||||
sub_configs: Dict[str, PretrainedConfig] = {
|
||||
"text_config": Qwen2Config(),
|
||||
"vision_config": SiglipVisionConfig(),
|
||||
}
|
||||
_auto_class: Optional[str] = "AutoConfig"
|
||||
|
||||
# Configuration for sub-modules.
|
||||
text_config: Qwen2Config = Qwen2Config()
|
||||
vision_config: SiglipVisionConfig = SiglipVisionConfig()
|
||||
|
||||
# Model configuration.
|
||||
hidden_size: int
|
||||
image_token_id: int
|
||||
mm_hidden_size: int
|
||||
mm_projector_type: str
|
||||
mm_vision_select_feature: str
|
||||
mm_vision_select_layer: int
|
||||
video_token_id: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_config: Optional[Dict[str, Any]] = None,
|
||||
vision_config: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
hidden_size: int = 1536,
|
||||
image_token_id: int = 151649,
|
||||
mm_hidden_size: int = 1152,
|
||||
mm_projector_type: str = "mlp_downsample_3x3_fix",
|
||||
mm_vision_select_feature: str = "cls_patch",
|
||||
mm_vision_select_layer: int = -2,
|
||||
video_token_id: int = 151650,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.text_config = Qwen2Config(**text_config) if text_config else Qwen2Config()
|
||||
self.vision_config = (
|
||||
SiglipVisionConfig(**vision_config)
|
||||
if vision_config
|
||||
else SiglipVisionConfig()
|
||||
)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.image_token_id = image_token_id
|
||||
self.mm_hidden_size = mm_hidden_size
|
||||
self.mm_projector_type = mm_projector_type
|
||||
self.mm_vision_select_feature = mm_vision_select_feature
|
||||
self.mm_vision_select_layer = mm_vision_select_layer
|
||||
self.video_token_id = video_token_id
|
||||
|
||||
|
||||
##### END COPY configuration.py #####
|
||||
|
||||
##### BEGIN COPY modeling_vila.py #####
|
||||
|
||||
|
||||
class DownSample3x3BlockFix(nn.Module):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
|
||||
|
||||
Returns:
|
||||
The output tensor of shape (batch_size, image_pad_len, mm_hidden_size * 9).
|
||||
"""
|
||||
|
||||
batch_size, sequence_length, hidden_size = x.shape
|
||||
|
||||
feat_size = int(sequence_length**0.5)
|
||||
if feat_size**2 != sequence_length:
|
||||
raise ValueError(
|
||||
f"Cannot take square root: sequence_length {sequence_length} is not a perfect square"
|
||||
)
|
||||
|
||||
features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
|
||||
|
||||
pad_after = (3 - feat_size % 3) % 3
|
||||
if pad_after > 0:
|
||||
features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
|
||||
feat_size = feat_size + pad_after
|
||||
|
||||
features = features.reshape(
|
||||
batch_size, feat_size // 3, 3, feat_size // 3, 3, hidden_size
|
||||
)
|
||||
features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
|
||||
features = features.reshape(batch_size, -1, 9 * hidden_size)
|
||||
|
||||
return features
|
||||
|
||||
|
||||
class MultimodalProjector(nn.Module):
|
||||
layers: nn.Sequential
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: VILAConfig,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
if config.mm_projector_type == "mlp_downsample_3x3_fix":
|
||||
self.layers = nn.Sequential(
|
||||
DownSample3x3BlockFix(),
|
||||
nn.LayerNorm(config.mm_hidden_size * 9),
|
||||
nn.Linear(
|
||||
config.mm_hidden_size * 9,
|
||||
config.mm_hidden_size * 3,
|
||||
),
|
||||
nn.GELU(),
|
||||
nn.LayerNorm(config.vision_config.hidden_size * 3),
|
||||
nn.Linear(config.vision_config.hidden_size * 3, config.hidden_size),
|
||||
nn.GELU(),
|
||||
nn.Linear(config.hidden_size, config.hidden_size),
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported mm_projector_type: {config.mm_projector_type}"
|
||||
)
|
||||
|
||||
self.layers.type(config.torch_dtype)
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: The input tensor of shape (batch_size, sequence_length, mm_hidden_size).
|
||||
|
||||
Returns:
|
||||
The output tensor of shape (batch_size, image_pad_len, hidden_size).
|
||||
"""
|
||||
|
||||
return self.layers(x.to(device=self.device, dtype=self.dtype))
|
||||
|
||||
|
||||
##### END COPY modeling_vila.py #####
|
||||
|
||||
|
||||
class VILAForConditionalGeneration(nn.Module):
|
||||
config: VILAConfig
|
||||
quant_config: Optional[QuantizationConfig]
|
||||
|
||||
logits_processor: LogitsProcessor
|
||||
pooler: Pooler
|
||||
|
||||
llm: Qwen2ForCausalLM
|
||||
mm_projector: MultimodalProjector
|
||||
vision_tower: SiglipVisionModel
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: VILAConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
|
||||
self.logits_processor = LogitsProcessor(config)
|
||||
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
|
||||
|
||||
self.llm = Qwen2ForCausalLM(
|
||||
config=config.text_config,
|
||||
quant_config=quant_config,
|
||||
prefix=utils.add_prefix("llm", prefix),
|
||||
)
|
||||
self.mm_projector = MultimodalProjector(config)
|
||||
self.vision_tower = SiglipVisionModel(config.vision_config)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return self.config.torch_dtype
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
positions: Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
get_embedding: bool = False,
|
||||
) -> LogitsProcessorOutput:
|
||||
output = mm_utils.general_mm_embed_routine(
|
||||
input_ids=input_ids,
|
||||
forward_batch=forward_batch,
|
||||
language_model=self.llm,
|
||||
data_embedding_funcs={
|
||||
Modality.IMAGE: self.get_image_feature,
|
||||
},
|
||||
get_embedding=get_embedding,
|
||||
positions=positions,
|
||||
)
|
||||
|
||||
return cast(LogitsProcessorOutput, output)
|
||||
|
||||
def get_image_feature(self, mm_input: List[MultimodalDataItem]) -> Tensor:
|
||||
pixel_values = cast(Tensor, mm_input[0].feature)
|
||||
|
||||
##### BEGIN COPY modeling_vila.py #####
|
||||
|
||||
vision_tower_output: BaseModelOutputWithPooling = self.vision_tower.__call__(
|
||||
pixel_values.to(
|
||||
device=self.vision_tower.device, dtype=self.vision_tower.dtype
|
||||
),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
|
||||
mm_projector_input = self._vision_tower_output_to_mm_projector_input(
|
||||
vision_tower_output
|
||||
)
|
||||
|
||||
image_embedding: Tensor = self.mm_projector.__call__(
|
||||
mm_projector_input.to(
|
||||
device=self.mm_projector.device, dtype=self.mm_projector.dtype
|
||||
)
|
||||
)
|
||||
|
||||
##### END COPY modeling_vila.py #####
|
||||
|
||||
return image_embedding
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, Tensor]]) -> None:
|
||||
params_dict = dict(self.named_parameters())
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if name.startswith("llm."):
|
||||
self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
|
||||
else:
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", weight_utils.default_weight_loader
|
||||
)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
def pad_input_ids(
|
||||
self, input_ids: List[int], mm_inputs: MultimodalInputs
|
||||
) -> List[int]:
|
||||
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||
return pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||
|
||||
##### BEGIN COPY modeling_vila.py #####
|
||||
|
||||
def _vision_tower_output_to_mm_projector_input(
|
||||
self,
|
||||
vision_tower_output: BaseModelOutputWithPooling,
|
||||
) -> Tensor:
|
||||
assert vision_tower_output.hidden_states is not None
|
||||
|
||||
selected_layer_hidden_states = vision_tower_output.hidden_states[
|
||||
self.config.mm_vision_select_layer
|
||||
]
|
||||
|
||||
if self.config.mm_vision_select_feature == "cls_patch":
|
||||
return selected_layer_hidden_states
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported mm_vision_select_feature: {self.config.mm_vision_select_feature}"
|
||||
)
|
||||
|
||||
##### END COPY modeling_vila.py #####
|
||||
|
||||
|
||||
EntryClass = [VILAForConditionalGeneration]
|
||||
@@ -185,6 +185,7 @@ class BaseMultimodalProcessor(ABC):
|
||||
"aspect_ratio_mask": Modality.IMAGE,
|
||||
"num_patches": Modality.IMAGE,
|
||||
"patch_pixel_values": Modality.IMAGE,
|
||||
"block_sizes": Modality.IMAGE,
|
||||
# Audio-related attributes
|
||||
"audio_features": Modality.AUDIO,
|
||||
"audio_feature_lens": Modality.AUDIO,
|
||||
|
||||
@@ -1,64 +1,72 @@
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import torch.nn as nn
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.processing_utils import ProcessorMixin
|
||||
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
|
||||
|
||||
from sglang.srt.managers.io_struct import (
|
||||
EmbeddingReqInput,
|
||||
GenerateReqInput,
|
||||
ImageDataInputItem,
|
||||
)
|
||||
from sglang.srt.models.vila import VILAForConditionalGeneration
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.srt.models.nvila import NVILAForConditionalGeneration
|
||||
from sglang.srt.models.nvila_lite import NVILALiteForConditionalGeneration
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
BaseMultimodalProcessor,
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
|
||||
class VILAProcessor(ProcessorMixin):
|
||||
"""A stub class for the VILA processor."""
|
||||
|
||||
tokenizer: PreTrainedTokenizerBase
|
||||
NUM_VIDEO_FRAMES = 8
|
||||
|
||||
|
||||
class VILAMultimodalProcessor(BaseMultimodalProcessor):
|
||||
models: List[Type[nn.Module]] = [VILAForConditionalGeneration]
|
||||
|
||||
_processor: VILAProcessor
|
||||
class NVILAMultimodalProcessor(BaseMultimodalProcessor):
|
||||
models: list[type[nn.Module]] = [
|
||||
NVILAForConditionalGeneration,
|
||||
NVILALiteForConditionalGeneration,
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hf_config: PretrainedConfig,
|
||||
server_args: ServerArgs,
|
||||
_processor: VILAProcessor,
|
||||
_processor: ProcessorMixin,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||
|
||||
self._processor: ProcessorMixin
|
||||
|
||||
tokenizer: PreTrainedTokenizerBase = getattr(self._processor, "tokenizer")
|
||||
|
||||
self.mm_tokens = MultimodalSpecialTokens(
|
||||
image_token=self._processor.tokenizer.image_token,
|
||||
image_token=tokenizer.image_token,
|
||||
image_token_id=hf_config.image_token_id,
|
||||
video_token=tokenizer.video_token,
|
||||
video_token_id=hf_config.video_token_id,
|
||||
).build(_processor)
|
||||
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
image_data: Optional[ImageDataInputItem | List[ImageDataInputItem]],
|
||||
input_text: str | List[int],
|
||||
request_obj: GenerateReqInput | EmbeddingReqInput,
|
||||
image_data,
|
||||
audio_data,
|
||||
input_text,
|
||||
request_obj: GenerateReqInput,
|
||||
**kwargs,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
) -> dict[str, Any] | None:
|
||||
base_output = self.load_mm_data(
|
||||
prompt=input_text,
|
||||
multimodal_tokens=self.mm_tokens,
|
||||
image_data=image_data,
|
||||
image_data=request_obj.image_data, # type: ignore
|
||||
video_data=request_obj.video_data, # type: ignore
|
||||
)
|
||||
|
||||
for i, video in enumerate(base_output.videos): # type: ignore
|
||||
base_output.videos[i] = [x.asnumpy() for x in video] # type: ignore
|
||||
|
||||
mm_items, input_ids, _ = self.process_and_combine_mm_data(
|
||||
base_output, self.mm_tokens
|
||||
base_output,
|
||||
self.mm_tokens,
|
||||
do_sample_frames=True,
|
||||
num_frames=NUM_VIDEO_FRAMES,
|
||||
)
|
||||
|
||||
return {
|
||||
Reference in New Issue
Block a user