447 lines
15 KiB
Python
447 lines
15 KiB
Python
# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
#
|
|
# Copyright 2025 The HuggingFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""SigLIP2 Vision Transformer for HunyuanImage3, rewritten in vLLM style.
|
|
|
|
Key optimizations over the original HuggingFace-style implementation:
|
|
- QKVParallelLinear: fused QKV projection with tensor parallelism support
|
|
- MMEncoderAttention: FlashAttention / xFormers backend
|
|
- ColumnParallelLinear / RowParallelLinear: TP-aware MLP layers
|
|
- Packed sequence processing: eliminates padding waste via cu_seqlens
|
|
- Data parallel support for multi-GPU ViT inference
|
|
"""
|
|
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.distributed import divide, get_tensor_model_parallel_world_size
|
|
from vllm.model_executor.layers.activation import get_act_fn
|
|
from vllm.model_executor.layers.attention import MMEncoderAttention
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.models.vision import is_vit_use_data_parallel
|
|
|
|
|
|
class Config:
|
|
"""Convert dict config to object with attribute access."""
|
|
|
|
def __init__(self, config):
|
|
if config is not None:
|
|
for key, value in config.items():
|
|
setattr(self, key, value)
|
|
|
|
def __getitem__(self, key):
|
|
return getattr(self, key, None)
|
|
|
|
def __setitem__(self, key, value):
|
|
return setattr(self, key, value)
|
|
|
|
|
|
class Siglip2VisionEmbeddings(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.embed_dim = config.hidden_size
|
|
self.patch_size = config.patch_size
|
|
|
|
self.patch_embedding = nn.Linear(
|
|
in_features=config.num_channels * self.patch_size * self.patch_size,
|
|
out_features=self.embed_dim,
|
|
)
|
|
|
|
self.num_patches = config.num_patches
|
|
self.position_embedding_size = int(self.num_patches**0.5)
|
|
self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
|
|
|
|
def forward(
|
|
self,
|
|
pixel_values: torch.FloatTensor,
|
|
spatial_shapes: torch.LongTensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Process packed pixel values with per-image position embeddings.
|
|
|
|
Args:
|
|
pixel_values: Packed pixel values
|
|
(total_real_patches, num_channels * patch_size * patch_size)
|
|
spatial_shapes: Per-image spatial shapes (B, 2) as [(h, w), ...]
|
|
|
|
Returns:
|
|
Packed embeddings (total_real_patches, embed_dim)
|
|
"""
|
|
target_dtype = self.patch_embedding.weight.dtype
|
|
patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
|
|
|
|
# Resize position embeddings per image and concatenate (packed)
|
|
positional_embeddings = self.position_embedding.weight.reshape(
|
|
self.position_embedding_size, self.position_embedding_size, -1
|
|
)
|
|
# (H_pe, W_pe, embed_dim) → (1, embed_dim, H_pe, W_pe) for interpolation
|
|
pe_for_resize = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
|
|
# Upcast on CPU since antialias is not supported for bfloat16/float16
|
|
if pe_for_resize.device.type == "cpu":
|
|
pe_for_resize = pe_for_resize.to(torch.float32)
|
|
|
|
position_embs: list[torch.Tensor] = []
|
|
for i in range(spatial_shapes.shape[0]):
|
|
height, width = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
|
|
resized = F.interpolate(
|
|
pe_for_resize,
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
antialias=True,
|
|
)
|
|
# (1, embed_dim, h, w) → (h*w, embed_dim)
|
|
resized = resized.reshape(self.embed_dim, height * width).transpose(0, 1)
|
|
position_embs.append(resized.to(target_dtype))
|
|
|
|
packed_position_embs = torch.cat(position_embs, dim=0)
|
|
return patch_embeds + packed_position_embs
|
|
|
|
|
|
class Siglip2Attention(nn.Module):
|
|
"""Multi-headed attention using QKVParallelLinear and MMEncoderAttention."""
|
|
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = config.hidden_size
|
|
self.num_heads = config.num_attention_heads
|
|
self.head_dim = self.embed_dim // self.num_heads
|
|
self.scale = self.head_dim**-0.5
|
|
|
|
use_data_parallel = is_vit_use_data_parallel()
|
|
self.qkv_proj = QKVParallelLinear(
|
|
hidden_size=self.embed_dim,
|
|
head_size=self.head_dim,
|
|
total_num_heads=self.num_heads,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.qkv_proj",
|
|
disable_tp=use_data_parallel,
|
|
)
|
|
self.out_proj = RowParallelLinear(
|
|
input_size=self.embed_dim,
|
|
output_size=self.embed_dim,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.out_proj",
|
|
disable_tp=use_data_parallel,
|
|
)
|
|
|
|
self.tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
|
|
self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
|
|
|
|
self.attn = MMEncoderAttention(
|
|
num_heads=self.num_heads_per_partition,
|
|
head_size=self.head_dim,
|
|
scale=self.scale,
|
|
prefix=f"{prefix}.attn",
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cu_seqlens: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
hidden_states: Packed input (total_tokens, embed_dim)
|
|
cu_seqlens: Cumulative sequence lengths (B+1,)
|
|
"""
|
|
seq_length = hidden_states.shape[0]
|
|
|
|
qkv, _ = self.qkv_proj(hidden_states)
|
|
q, k, v = qkv.chunk(3, dim=-1)
|
|
|
|
q = q.view(seq_length, self.num_heads_per_partition, self.head_dim)
|
|
k = k.view(seq_length, self.num_heads_per_partition, self.head_dim)
|
|
v = v.view(seq_length, self.num_heads_per_partition, self.head_dim)
|
|
|
|
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
|
|
attn_output = self.attn(
|
|
query=q.unsqueeze(0),
|
|
key=k.unsqueeze(0),
|
|
value=v.unsqueeze(0),
|
|
cu_seqlens=cu_seqlens,
|
|
max_seqlen=max_seqlen,
|
|
)
|
|
attn_output = attn_output.reshape(seq_length, self.num_heads_per_partition * self.head_dim)
|
|
|
|
attn_output, _ = self.out_proj(attn_output)
|
|
return attn_output
|
|
|
|
|
|
class Siglip2MLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
use_data_parallel = is_vit_use_data_parallel()
|
|
self.activation_fn = get_act_fn(config.hidden_act)
|
|
self.fc1 = ColumnParallelLinear(
|
|
config.hidden_size,
|
|
config.intermediate_size,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.fc1",
|
|
disable_tp=use_data_parallel,
|
|
)
|
|
self.fc2 = RowParallelLinear(
|
|
config.intermediate_size,
|
|
config.hidden_size,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.fc2",
|
|
disable_tp=use_data_parallel,
|
|
)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states, _ = self.fc1(hidden_states)
|
|
hidden_states = self.activation_fn(hidden_states)
|
|
hidden_states, _ = self.fc2(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class Siglip2EncoderLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = config.hidden_size
|
|
self.self_attn = Siglip2Attention(
|
|
config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.self_attn",
|
|
)
|
|
self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
|
|
self.mlp = Siglip2MLP(
|
|
config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.mlp",
|
|
)
|
|
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cu_seqlens: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
hidden_states: Packed input (total_tokens, embed_dim)
|
|
cu_seqlens: Cumulative sequence lengths (B+1,)
|
|
"""
|
|
residual = hidden_states
|
|
hidden_states = self.layer_norm1(hidden_states)
|
|
hidden_states = self.self_attn(hidden_states, cu_seqlens=cu_seqlens)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.layer_norm2(hidden_states)
|
|
hidden_states = self.mlp(hidden_states)
|
|
hidden_states = residual + hidden_states
|
|
return hidden_states
|
|
|
|
|
|
class Siglip2Encoder(nn.Module):
|
|
"""Transformer encoder with packed sequence processing."""
|
|
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
Siglip2EncoderLayer(
|
|
config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.layers.{idx}",
|
|
)
|
|
for idx in range(config.num_hidden_layers)
|
|
]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cu_seqlens: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
hidden_states: Packed input (total_tokens, embed_dim)
|
|
cu_seqlens: Cumulative sequence lengths (B+1,)
|
|
"""
|
|
for layer in self.layers:
|
|
hidden_states = layer(hidden_states, cu_seqlens)
|
|
return hidden_states
|
|
|
|
|
|
class Siglip2VisionTransformer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
config = Config(config)
|
|
self.config = config
|
|
self.embed_dim = config.hidden_size
|
|
|
|
self.embeddings = Siglip2VisionEmbeddings(config)
|
|
self.encoder = Siglip2Encoder(
|
|
config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.encoder" if prefix else "encoder",
|
|
)
|
|
self.post_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
pixel_values: torch.FloatTensor,
|
|
attention_mask: torch.Tensor,
|
|
spatial_shapes: torch.LongTensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
pixel_values: Batched pixel values
|
|
(B, max_num_patches, num_channels * patch_size * patch_size)
|
|
attention_mask: (B, max_num_patches) with 1 for real, 0 for padding
|
|
spatial_shapes: (B, 2) with (height, width) per image
|
|
|
|
Returns:
|
|
(B, max_num_patches, hidden_size) with zeros at padding positions
|
|
"""
|
|
batch_size, max_patches, _ = pixel_values.shape
|
|
|
|
# Pack: extract real tokens using attention_mask
|
|
mask_bool = attention_mask.bool()
|
|
packed_pixels = pixel_values[mask_bool] # (total_real_patches, C*P*P)
|
|
|
|
# Compute cu_seqlens from spatial_shapes (h * w per image)
|
|
seq_lens = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(torch.int32)
|
|
cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=pixel_values.device)
|
|
cu_seqlens[1:] = seq_lens.cumsum(0)
|
|
|
|
# Embeddings (packed)
|
|
hidden_states = self.embeddings(packed_pixels, spatial_shapes)
|
|
|
|
# Encoder (packed)
|
|
hidden_states = self.encoder(hidden_states, cu_seqlens)
|
|
|
|
# Post layernorm
|
|
hidden_states = self.post_layernorm(hidden_states)
|
|
|
|
# Unpack: scatter back to (B, max_patches, hidden_size)
|
|
output = torch.zeros(
|
|
batch_size,
|
|
max_patches,
|
|
self.embed_dim,
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
output[mask_bool] = hidden_states
|
|
|
|
return output
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
stacked_params_mapping = [
|
|
# (param_name, shard_name, shard_id)
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
]
|
|
params_dict = dict(self.named_parameters())
|
|
loaded_params: set[str] = set()
|
|
|
|
for name, loaded_weight in weights:
|
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
|
if weight_name not in name:
|
|
continue
|
|
name = name.replace(weight_name, param_name)
|
|
param = params_dict[name]
|
|
weight_loader = param.weight_loader
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
break
|
|
else:
|
|
if name not in params_dict:
|
|
# Skip weights for removed modules (e.g., pooling head)
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(name)
|
|
return loaded_params
|
|
|
|
|
|
class LightProjector(nn.Module):
|
|
def __init__(self, config):
|
|
config = Config(config)
|
|
super().__init__()
|
|
|
|
if config.projector_type == "linear":
|
|
modules = nn.Linear(config.input_dim, config.n_embed)
|
|
|
|
elif config.projector_type == "mlp_gelu":
|
|
modules = [nn.Linear(config.input_dim, config.n_embed)]
|
|
for _ in range(1, config.depth):
|
|
modules.append(nn.GELU())
|
|
modules.append(nn.Linear(config.n_embed, config.n_embed))
|
|
modules = nn.Sequential(*modules)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown projector type: {config.projector_type}")
|
|
|
|
self.layers = modules
|
|
|
|
def forward(self, x):
|
|
return self.layers(x)
|
|
|
|
|
|
__all__ = [
|
|
"Siglip2VisionTransformer",
|
|
"LightProjector",
|
|
]
|