1366 lines
49 KiB
Python
1366 lines
49 KiB
Python
import itertools
|
|
import logging
|
|
from collections.abc import Iterable
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import vllm._custom_ops as ops
|
|
from diffusers.models.activations import get_activation
|
|
from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
from einops import rearrange, repeat
|
|
from vllm.model_executor.layers.activation import get_act_and_mul_fn
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.linear import (
|
|
MergedColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
from vllm_omni.diffusion.distributed.sp_plan import (
|
|
SequenceParallelInput,
|
|
SequenceParallelOutput,
|
|
)
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _patch_cutlass_padded_fp8():
|
|
"""Monkey-patch vllm._custom_ops.cutlass_scaled_mm to pad tensors whose
|
|
dimensions are not multiples of 16, so the CUTLASS FP8 kernel is used.
|
|
|
|
OmniGen2 has hidden_size=2520 (2520 % 16 == 8). Without this patch,
|
|
vLLM's cutlass_scaled_mm falls back to a Triton scaled_mm kernel for
|
|
every FP8 linear layer (QKV, attn output, gate_up_proj, down_proj),
|
|
which is dramatically slower than the native CUTLASS FP8 tensor-core
|
|
path on H100/H200 GPUs.
|
|
|
|
Weight tensors (b) are constant across forward passes, so padded
|
|
versions are computed once and cached by data_ptr to avoid repeated
|
|
allocation and column-major conversion overhead.
|
|
"""
|
|
_orig_cutlass_scaled_mm = ops.cutlass_scaled_mm
|
|
# Cache: data_ptr → (padded_b, padded_bias, padded_scale_b, pad_k, pad_n, orig_n)
|
|
_weight_cache: dict[int, tuple] = {}
|
|
|
|
def _padded_cutlass_scaled_mm(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
scale_a: torch.Tensor,
|
|
scale_b: torch.Tensor,
|
|
out_dtype: torch.dtype,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0:
|
|
return _orig_cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
|
|
|
# Reshape to 2D (mirrors the original function)
|
|
target_shape = (*a.shape[:-1], b.shape[1])
|
|
a = a.view(-1, a.shape[-1])
|
|
orig_n = b.shape[1]
|
|
|
|
# Cache the padded weight — it's a model parameter that never changes.
|
|
key = b.data_ptr()
|
|
if key not in _weight_cache:
|
|
pad_k = (16 - b.shape[0] % 16) % 16
|
|
pad_n = (16 - orig_n % 16) % 16
|
|
b_pad = b
|
|
if pad_k > 0:
|
|
b_pad = F.pad(b_pad, (0, 0, 0, pad_k))
|
|
if pad_n > 0:
|
|
b_pad = F.pad(b_pad, (0, pad_n))
|
|
# CUTLASS requires b column-major (stride(0)==1).
|
|
b_pad = b_pad.t().contiguous().t()
|
|
|
|
bias_pad = None
|
|
if bias is not None and pad_n > 0:
|
|
bias_pad = F.pad(bias, (0, pad_n))
|
|
|
|
scale_b_pad = scale_b
|
|
if scale_b.numel() > 1 and pad_n > 0:
|
|
scale_b_pad = F.pad(
|
|
scale_b.view(-1, scale_b.shape[-1]),
|
|
(0, pad_n),
|
|
value=1.0,
|
|
)
|
|
|
|
_weight_cache[key] = (
|
|
b_pad,
|
|
bias_pad,
|
|
scale_b_pad,
|
|
pad_k,
|
|
pad_n,
|
|
orig_n,
|
|
)
|
|
|
|
b_pad, bias_pad, scale_b_pad, pad_k, pad_n, orig_n = _weight_cache[key]
|
|
|
|
# Pad activations on K dimension (cheap — activations are small).
|
|
if pad_k > 0:
|
|
a = F.pad(a, (0, pad_k)).contiguous()
|
|
|
|
out = torch.empty((a.shape[0], b_pad.shape[1]), dtype=out_dtype, device=a.device)
|
|
torch.ops._C.cutlass_scaled_mm(
|
|
out,
|
|
a,
|
|
b_pad,
|
|
scale_a,
|
|
scale_b_pad,
|
|
bias_pad if bias is not None else None,
|
|
)
|
|
|
|
if pad_n > 0:
|
|
out = out[:, :orig_n]
|
|
|
|
return out.view(*target_shape)
|
|
|
|
ops.cutlass_scaled_mm = _padded_cutlass_scaled_mm
|
|
logger.info(
|
|
"Patched vllm._custom_ops.cutlass_scaled_mm with CUTLASS-padded FP8 "
|
|
"variant (avoids slow Triton fallback for non-%%16 dimensions)"
|
|
)
|
|
|
|
|
|
_patch_cutlass_padded_fp8()
|
|
|
|
|
|
class OmniGen2Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
num_kv_heads: int,
|
|
eps: float = 1e-5,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.num_kv_heads = num_kv_heads
|
|
self.head_dim = dim // num_heads
|
|
self.eps = eps
|
|
|
|
self.to_qkv = QKVParallelLinear(
|
|
hidden_size=dim,
|
|
head_size=self.head_dim,
|
|
total_num_heads=num_heads,
|
|
total_num_kv_heads=num_kv_heads,
|
|
disable_tp=True,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.to_qkv",
|
|
)
|
|
|
|
self.norm_q = RMSNorm(self.head_dim, eps=eps)
|
|
self.norm_k = RMSNorm(self.head_dim, eps=eps)
|
|
|
|
self.to_out = nn.ModuleList(
|
|
[
|
|
RowParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=False,
|
|
input_is_parallel=False,
|
|
quant_config=quant_config,
|
|
return_bias=False,
|
|
prefix=f"{prefix}.to_out.0",
|
|
)
|
|
]
|
|
)
|
|
self.attn = Attention(
|
|
num_heads=num_heads,
|
|
head_size=self.head_dim,
|
|
softmax_scale=1.0 / (self.head_dim**0.5),
|
|
causal=False,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: torch.Tensor | None = None,
|
|
image_rotary_emb: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Process attention computation with flash attention.
|
|
|
|
Args:
|
|
hidden_states: Hidden states tensor of shape (batch_size, seq_len, hidden_dim)
|
|
attention_mask: Optional attention mask tensor
|
|
image_rotary_emb: Optional rotary embeddings for image tokens
|
|
|
|
Returns:
|
|
torch.Tensor: Processed hidden states after attention computation
|
|
"""
|
|
batch_size = hidden_states.shape[0]
|
|
|
|
# Contiguous layout for FP8 quantized linear GEMMs (matches FLUX DiT).
|
|
hidden_states = hidden_states.contiguous()
|
|
|
|
# Get Query-Key-Value Pair
|
|
qkv, _ = self.to_qkv(hidden_states)
|
|
|
|
q_size = self.num_heads * self.head_dim
|
|
kv_size = self.num_kv_heads * self.head_dim
|
|
query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
|
|
|
# Reshape tensors for attention computation
|
|
query = query.unflatten(-1, (self.num_heads, -1))
|
|
key = key.unflatten(-1, (self.num_kv_heads, -1))
|
|
value = value.unflatten(-1, (self.num_kv_heads, -1))
|
|
|
|
# Apply Query-Key normalization
|
|
if self.norm_q is not None:
|
|
query = self.norm_q(query)
|
|
if self.norm_k is not None:
|
|
key = self.norm_k(key)
|
|
|
|
dtype = query.dtype
|
|
|
|
# Apply Rotary Position Embeddings
|
|
if image_rotary_emb is not None:
|
|
query = apply_rotary_emb(query, image_rotary_emb, use_real=False)
|
|
key = apply_rotary_emb(key, image_rotary_emb, use_real=False)
|
|
|
|
query, key = query.to(dtype), key.to(dtype)
|
|
|
|
if self.num_kv_heads < self.num_heads:
|
|
# Repeat key and value heads to match query heads
|
|
num_groups = self.num_heads // self.num_kv_heads
|
|
key = key.repeat_interleave(num_groups, dim=2)
|
|
value = value.repeat_interleave(num_groups, dim=2)
|
|
|
|
hidden_states = self.attn(
|
|
query,
|
|
key,
|
|
value,
|
|
)
|
|
|
|
# Reshape back
|
|
hidden_states = hidden_states.reshape(batch_size, -1, self.num_heads * self.head_dim)
|
|
hidden_states = hidden_states.to(dtype)
|
|
|
|
hidden_states = self.to_out[0](hidden_states.contiguous())
|
|
|
|
return hidden_states
|
|
|
|
|
|
class TimestepEmbedding(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
time_embed_dim: int,
|
|
act_fn: str = "silu",
|
|
out_dim: int = None,
|
|
post_act_fn: str | None = None,
|
|
cond_proj_dim=None,
|
|
sample_proj_bias=True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
|
|
|
|
if cond_proj_dim is not None:
|
|
self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
|
|
else:
|
|
self.cond_proj = None
|
|
|
|
self.act = get_activation(act_fn)
|
|
|
|
if out_dim is not None:
|
|
time_embed_dim_out = out_dim
|
|
else:
|
|
time_embed_dim_out = time_embed_dim
|
|
self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
|
|
|
|
if post_act_fn is None:
|
|
self.post_act = None
|
|
else:
|
|
self.post_act = get_activation(post_act_fn)
|
|
|
|
def forward(self, sample, condition=None):
|
|
if condition is not None:
|
|
sample = sample + self.cond_proj(condition)
|
|
sample = self.linear_1(sample)
|
|
sample = self.act(sample)
|
|
sample = self.linear_2(sample)
|
|
|
|
if self.post_act is not None:
|
|
sample = self.post_act(sample)
|
|
return sample
|
|
|
|
|
|
def apply_rotary_emb(
|
|
x: torch.Tensor,
|
|
freqs_cis: torch.Tensor | tuple[torch.Tensor],
|
|
use_real: bool = True,
|
|
use_real_unbind_dim: int = -1,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
|
|
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
|
|
reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
|
|
tensors contain rotary embeddings and are returned as real tensors.
|
|
|
|
Args:
|
|
x (`torch.Tensor`):
|
|
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
|
|
freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
|
|
"""
|
|
if use_real:
|
|
cos, sin = freqs_cis # [S, D]
|
|
cos = cos[None, None]
|
|
sin = sin[None, None]
|
|
cos, sin = cos.to(x.device), sin.to(x.device)
|
|
|
|
if use_real_unbind_dim == -1:
|
|
# Used for flux, cogvideox, hunyuan-dit
|
|
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
|
|
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
|
elif use_real_unbind_dim == -2:
|
|
# Used for Stable Audio, OmniGen and CogView4
|
|
x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2]
|
|
x_rotated = torch.cat([-x_imag, x_real], dim=-1)
|
|
else:
|
|
raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
|
|
|
|
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
|
|
|
|
return out
|
|
else:
|
|
# used for lumina
|
|
x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2))
|
|
freqs_cis = freqs_cis.unsqueeze(2)
|
|
x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
|
|
|
|
return x_out.type_as(x)
|
|
|
|
|
|
class LuminaRMSNormZero(nn.Module):
|
|
"""
|
|
Norm layer adaptive RMS normalization zero.
|
|
|
|
Parameters:
|
|
embedding_dim (`int`): The size of each embedding vector.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
embedding_dim: int,
|
|
norm_eps: float,
|
|
norm_elementwise_affine: bool,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(
|
|
min(embedding_dim, 1024),
|
|
4 * embedding_dim,
|
|
bias=True,
|
|
)
|
|
|
|
self.norm = RMSNorm(embedding_dim, eps=norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
emb: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
emb = self.linear(self.silu(emb))
|
|
scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
|
|
x = self.norm(x) * (1 + scale_msa[:, None])
|
|
return x, gate_msa, scale_mlp, gate_mlp
|
|
|
|
|
|
class LuminaLayerNormContinuous(nn.Module):
|
|
def __init__(
|
|
self,
|
|
embedding_dim: int,
|
|
conditioning_embedding_dim: int,
|
|
# NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
|
|
# because the output is immediately scaled and shifted by the projected conditioning embeddings.
|
|
# Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
|
|
# However, this is how it was implemented in the original code, and it's rather likely you should
|
|
# set `elementwise_affine` to False.
|
|
elementwise_affine=True,
|
|
eps=1e-5,
|
|
bias=True,
|
|
norm_type="layer_norm",
|
|
out_dim: int | None = None,
|
|
):
|
|
super().__init__()
|
|
|
|
# AdaLN
|
|
self.silu = nn.SiLU()
|
|
self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
|
|
|
|
if norm_type == "layer_norm":
|
|
self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
|
|
elif norm_type == "rms_norm":
|
|
self.norm = RMSNorm(embedding_dim, eps=eps)
|
|
else:
|
|
raise ValueError(f"unknown norm_type {norm_type}")
|
|
|
|
self.linear_2 = None
|
|
if out_dim is not None:
|
|
self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
conditioning_embedding: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
# convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32
|
|
# (needed for hunyuanDiT)
|
|
emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
|
|
scale = emb
|
|
x = self.norm(x) * (1 + scale)[:, None, :]
|
|
|
|
if self.linear_2 is not None:
|
|
x = self.linear_2(x)
|
|
|
|
return x
|
|
|
|
|
|
class LuminaFeedForward(nn.Module):
|
|
r"""
|
|
A feed-forward layer.
|
|
|
|
Parameters:
|
|
dim (`int`):
|
|
The dimensionality of the input and output tensors.
|
|
inner_dim (`int`): The intermediate dimension of the feedforward layer.
|
|
multiple_of (`int`, *optional*): Value to ensure hidden dimension is a multiple
|
|
of this value.
|
|
ffn_dim_multiplier (float, *optional*): Custom multiplier for hidden
|
|
dimension. Defaults to None.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
inner_dim: int,
|
|
multiple_of: int | None = 256,
|
|
ffn_dim_multiplier: float | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
|
|
# custom hidden_size factor multiplier
|
|
if ffn_dim_multiplier is not None:
|
|
inner_dim = int(ffn_dim_multiplier * inner_dim)
|
|
inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
|
|
|
|
self.gate_up_proj = MergedColumnParallelLinear(
|
|
dim,
|
|
[inner_dim, inner_dim],
|
|
bias=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.gate_up_proj",
|
|
)
|
|
self.act_fn = get_act_and_mul_fn("silu")
|
|
self.down_proj = RowParallelLinear(
|
|
inner_dim,
|
|
dim,
|
|
bias=False,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.down_proj",
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.gate_up_proj(x)
|
|
x = self.act_fn(x)
|
|
return self.down_proj(x)
|
|
|
|
|
|
class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int = 4096,
|
|
text_feat_dim: int = 2048,
|
|
frequency_embedding_size: int = 256,
|
|
norm_eps: float = 1e-5,
|
|
timestep_scale: float = 1.0,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.time_proj = Timesteps(
|
|
num_channels=frequency_embedding_size,
|
|
flip_sin_to_cos=True,
|
|
downscale_freq_shift=0.0,
|
|
scale=timestep_scale,
|
|
)
|
|
|
|
self.timestep_embedder = TimestepEmbedding(
|
|
in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
|
|
)
|
|
|
|
self.caption_embedder = nn.Sequential(
|
|
RMSNorm(text_feat_dim, eps=norm_eps),
|
|
nn.Linear(text_feat_dim, hidden_size, bias=True),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
timestep: torch.Tensor,
|
|
text_hidden_states: torch.Tensor,
|
|
dtype: torch.dtype,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
timestep_proj = self.time_proj(timestep).to(dtype=dtype)
|
|
time_embed = self.timestep_embedder(timestep_proj)
|
|
caption_embed = self.caption_embedder(text_hidden_states)
|
|
return time_embed, caption_embed
|
|
|
|
|
|
class OmniGen2RotaryPosEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
theta: int,
|
|
axes_dim: tuple[int, int, int],
|
|
axes_lens: tuple[int, int, int] = (300, 512, 512),
|
|
patch_size: int = 2,
|
|
):
|
|
super().__init__()
|
|
self.theta = theta
|
|
self.axes_dim = axes_dim
|
|
self.axes_lens = axes_lens
|
|
self.patch_size = patch_size
|
|
|
|
@staticmethod
|
|
def get_freqs_cis(
|
|
axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
|
|
) -> list[torch.Tensor]:
|
|
freqs_cis = []
|
|
freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
|
|
for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
|
|
emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
|
|
freqs_cis.append(emb)
|
|
return freqs_cis
|
|
|
|
def _get_freqs_cis(self, freqs_cis, ids: torch.Tensor) -> torch.Tensor:
|
|
device = ids.device
|
|
if ids.device.type == "mps":
|
|
ids = ids.to("cpu")
|
|
|
|
result = []
|
|
for i in range(len(self.axes_dim)):
|
|
freqs = freqs_cis[i].to(ids.device)
|
|
index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64)
|
|
result.append(torch.gather(freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
|
|
return torch.cat(result, dim=-1).to(device)
|
|
|
|
def forward(
|
|
self,
|
|
freqs_cis,
|
|
attention_mask,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
ref_img_sizes,
|
|
img_sizes,
|
|
device,
|
|
):
|
|
batch_size = len(attention_mask)
|
|
p = self.patch_size
|
|
|
|
encoder_seq_len = attention_mask.shape[1]
|
|
l_effective_cap_len = attention_mask.sum(dim=1).tolist()
|
|
|
|
seq_lengths = [
|
|
cap_len + sum(ref_img_len) + img_len
|
|
for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len)
|
|
]
|
|
|
|
max_seq_len = max(seq_lengths)
|
|
max_ref_img_len = max([sum(ref_img_len) for ref_img_len in l_effective_ref_img_len])
|
|
max_img_len = max(l_effective_img_len)
|
|
|
|
# Create position IDs
|
|
position_ids = torch.zeros(batch_size, max_seq_len, 3, dtype=torch.int32, device=device)
|
|
|
|
for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)):
|
|
# add text position ids
|
|
position_ids[i, :cap_seq_len] = repeat(
|
|
torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
|
|
)
|
|
|
|
pe_shift = cap_seq_len
|
|
pe_shift_len = cap_seq_len
|
|
|
|
if ref_img_sizes[i] is not None:
|
|
for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]):
|
|
H, W = ref_img_size
|
|
ref_H_tokens, ref_W_tokens = H // p, W // p
|
|
assert ref_H_tokens * ref_W_tokens == ref_img_len
|
|
# add image position ids
|
|
|
|
row_ids = repeat(
|
|
torch.arange(ref_H_tokens, dtype=torch.int32, device=device),
|
|
"h -> h w",
|
|
w=ref_W_tokens,
|
|
).flatten()
|
|
col_ids = repeat(
|
|
torch.arange(ref_W_tokens, dtype=torch.int32, device=device),
|
|
"w -> h w",
|
|
h=ref_H_tokens,
|
|
).flatten()
|
|
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift
|
|
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids
|
|
position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids
|
|
|
|
pe_shift += max(ref_H_tokens, ref_W_tokens)
|
|
pe_shift_len += ref_img_len
|
|
|
|
H, W = img_sizes[i]
|
|
H_tokens, W_tokens = H // p, W // p
|
|
assert H_tokens * W_tokens == l_effective_img_len[i]
|
|
|
|
row_ids = repeat(
|
|
torch.arange(H_tokens, dtype=torch.int32, device=device),
|
|
"h -> h w",
|
|
w=W_tokens,
|
|
).flatten()
|
|
col_ids = repeat(
|
|
torch.arange(W_tokens, dtype=torch.int32, device=device),
|
|
"w -> h w",
|
|
h=H_tokens,
|
|
).flatten()
|
|
|
|
assert pe_shift_len + l_effective_img_len[i] == seq_len
|
|
position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
|
|
position_ids[i, pe_shift_len:seq_len, 1] = row_ids
|
|
position_ids[i, pe_shift_len:seq_len, 2] = col_ids
|
|
|
|
# Get combined rotary embeddings
|
|
freqs_cis = self._get_freqs_cis(freqs_cis, position_ids)
|
|
|
|
# create separate rotary embeddings for captions and images
|
|
cap_freqs_cis = torch.zeros(
|
|
batch_size,
|
|
encoder_seq_len,
|
|
freqs_cis.shape[-1],
|
|
device=device,
|
|
dtype=freqs_cis.dtype,
|
|
)
|
|
ref_img_freqs_cis = torch.zeros(
|
|
batch_size,
|
|
max_ref_img_len,
|
|
freqs_cis.shape[-1],
|
|
device=device,
|
|
dtype=freqs_cis.dtype,
|
|
)
|
|
img_freqs_cis = torch.zeros(
|
|
batch_size,
|
|
max_img_len,
|
|
freqs_cis.shape[-1],
|
|
device=device,
|
|
dtype=freqs_cis.dtype,
|
|
)
|
|
|
|
for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
|
|
zip(
|
|
l_effective_cap_len,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
seq_lengths,
|
|
)
|
|
):
|
|
cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len]
|
|
ref_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
|
|
img_freqs_cis[i, :img_len] = freqs_cis[
|
|
i,
|
|
cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len,
|
|
]
|
|
|
|
return (
|
|
cap_freqs_cis,
|
|
ref_img_freqs_cis,
|
|
img_freqs_cis,
|
|
freqs_cis,
|
|
l_effective_cap_len,
|
|
seq_lengths,
|
|
)
|
|
|
|
|
|
class OmniGen2TransformerBlock(nn.Module):
|
|
"""
|
|
Transformer block for OmniGen2 model.
|
|
|
|
This block implements a transformer layer with:
|
|
- Multi-head attention with flash attention
|
|
- Feed-forward network with SwiGLU activation
|
|
- RMS normalization
|
|
- Optional modulation for conditional generation
|
|
|
|
Args:
|
|
dim: Dimension of the input and output tensors
|
|
num_attention_heads: Number of attention heads
|
|
num_kv_heads: Number of key-value heads
|
|
multiple_of: Multiple of which the hidden dimension should be
|
|
ffn_dim_multiplier: Multiplier for the feed-forward network dimension
|
|
norm_eps: Epsilon value for normalization layers
|
|
modulation: Whether to use modulation for conditional generation
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_attention_heads: int,
|
|
num_kv_heads: int,
|
|
multiple_of: int,
|
|
ffn_dim_multiplier: float,
|
|
norm_eps: float,
|
|
modulation: bool = True,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
"""Initialize the transformer block."""
|
|
super().__init__()
|
|
self.head_dim = dim // num_attention_heads
|
|
self.modulation = modulation
|
|
|
|
self.attn = OmniGen2Attention(
|
|
dim=dim,
|
|
num_heads=num_attention_heads,
|
|
num_kv_heads=num_kv_heads,
|
|
eps=1e-5,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.attn",
|
|
)
|
|
|
|
# Initialize feed-forward network
|
|
self.feed_forward = LuminaFeedForward(
|
|
dim=dim,
|
|
inner_dim=4 * dim,
|
|
multiple_of=multiple_of,
|
|
ffn_dim_multiplier=ffn_dim_multiplier,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.feed_forward",
|
|
)
|
|
|
|
# Initialize normalization layers
|
|
if modulation:
|
|
self.norm1 = LuminaRMSNormZero(
|
|
embedding_dim=dim,
|
|
norm_eps=norm_eps,
|
|
norm_elementwise_affine=True,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.norm1",
|
|
)
|
|
else:
|
|
self.norm1 = RMSNorm(dim, eps=norm_eps)
|
|
|
|
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
|
|
self.norm2 = RMSNorm(dim, eps=norm_eps)
|
|
self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
image_rotary_emb: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Forward pass of the transformer block.
|
|
|
|
Args:
|
|
hidden_states: Input hidden states tensor
|
|
attention_mask: Attention mask tensor
|
|
image_rotary_emb: Rotary embeddings for image tokens
|
|
temb: Optional timestep embedding tensor
|
|
|
|
Returns:
|
|
torch.Tensor: Output hidden states after transformer block processing
|
|
"""
|
|
if self.modulation:
|
|
if temb is None:
|
|
raise ValueError("temb must be provided when modulation is enabled")
|
|
|
|
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
|
|
attn_output = self.attn(
|
|
hidden_states=norm_hidden_states,
|
|
attention_mask=attention_mask,
|
|
image_rotary_emb=image_rotary_emb,
|
|
)
|
|
hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output)
|
|
mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
|
|
hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
|
|
else:
|
|
norm_hidden_states = self.norm1(hidden_states)
|
|
attn_output = self.attn(
|
|
hidden_states=norm_hidden_states,
|
|
attention_mask=attention_mask,
|
|
image_rotary_emb=image_rotary_emb,
|
|
)
|
|
hidden_states = hidden_states + self.norm2(attn_output)
|
|
mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
|
|
hidden_states = hidden_states + self.ffn_norm2(mlp_output)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class _OmniGen2SPInputBoundary(nn.Module):
|
|
"""SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.
|
|
|
|
Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer
|
|
once per CFG branch and the unconditional / conditional branches have
|
|
different joint sequence lengths (text prefix differs). The SP framework's
|
|
auto-pad records the length only on the first call, so without this reset
|
|
a stale length from an earlier branch truncates the gather of later ones.
|
|
"""
|
|
|
|
def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
from vllm_omni.diffusion.forward_context import (
|
|
get_forward_context,
|
|
is_forward_context_available,
|
|
)
|
|
|
|
if is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
ctx.sp_original_seq_len = None
|
|
ctx.sp_padding_size = 0
|
|
return hidden_states, rotary_emb
|
|
|
|
|
|
class _OmniGen2SPOutputBoundary(nn.Module):
|
|
"""SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``."""
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
return hidden_states
|
|
|
|
|
|
class OmniGen2Transformer2DModel(nn.Module):
|
|
"""
|
|
OmniGen2 Transformer 2D Model.
|
|
|
|
A transformer-based diffusion model for image generation with:
|
|
- Patch-based image processing
|
|
- Rotary position embeddings
|
|
- Multi-head attention
|
|
- Conditional generation support
|
|
|
|
Args:
|
|
patch_size: Size of image patches
|
|
in_channels: Number of input channels
|
|
out_channels: Number of output channels (defaults to in_channels)
|
|
hidden_size: Size of hidden layers
|
|
num_layers: Number of transformer layers
|
|
num_refiner_layers: Number of refiner layers
|
|
num_attention_heads: Number of attention heads
|
|
num_kv_heads: Number of key-value heads
|
|
multiple_of: Multiple of which the hidden dimension should be
|
|
ffn_dim_multiplier: Multiplier for feed-forward network dimension
|
|
norm_eps: Epsilon value for normalization layers
|
|
axes_dim_rope: Dimensions for rotary position embeddings
|
|
axes_lens: Lengths for rotary position embeddings
|
|
text_feat_dim: Dimension of text features
|
|
timestep_scale: Scale factor for timestep embeddings
|
|
"""
|
|
|
|
# ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention
|
|
# accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP
|
|
# cannot use attention masks.
|
|
_sp_plan = {
|
|
"sp_input_boundary": {
|
|
0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
|
|
1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
|
|
},
|
|
"sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
patch_size: int = 2,
|
|
in_channels: int = 16,
|
|
out_channels: int | None = None,
|
|
hidden_size: int = 2520,
|
|
num_layers: int = 32,
|
|
num_refiner_layers: int = 2,
|
|
num_attention_heads: int = 21,
|
|
num_kv_heads: int = 7,
|
|
multiple_of: int = 256,
|
|
ffn_dim_multiplier: float | None = None,
|
|
norm_eps: float = 1e-5,
|
|
axes_dim_rope: tuple[int, int, int] = (40, 40, 40),
|
|
axes_lens: tuple[int, int, int] = (1024, 1664, 1664),
|
|
text_feat_dim: int = 2048,
|
|
timestep_scale: float = 1000.0,
|
|
quant_config: QuantizationConfig | None = None,
|
|
) -> None:
|
|
"""Initialize the OmniGen2 transformer model."""
|
|
super().__init__()
|
|
|
|
# Validate configuration
|
|
if (hidden_size // num_attention_heads) != sum(axes_dim_rope):
|
|
raise ValueError(
|
|
f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) "
|
|
f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})"
|
|
)
|
|
|
|
self.config = SimpleNamespace(
|
|
patch_size=patch_size,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels or in_channels,
|
|
hidden_size=hidden_size,
|
|
axes_dim_rope=axes_dim_rope,
|
|
axes_lens=axes_lens,
|
|
)
|
|
self.out_channels = self.config.out_channels
|
|
|
|
# Initialize embeddings
|
|
self.rope_embedder = OmniGen2RotaryPosEmbed(
|
|
theta=10000,
|
|
axes_dim=axes_dim_rope,
|
|
axes_lens=axes_lens,
|
|
patch_size=patch_size,
|
|
)
|
|
|
|
self.x_embedder = nn.Linear(
|
|
in_features=patch_size * patch_size * in_channels,
|
|
out_features=hidden_size,
|
|
)
|
|
|
|
self.ref_image_patch_embedder = nn.Linear(
|
|
in_features=patch_size * patch_size * in_channels,
|
|
out_features=hidden_size,
|
|
)
|
|
|
|
self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
|
|
hidden_size=hidden_size,
|
|
text_feat_dim=text_feat_dim,
|
|
norm_eps=norm_eps,
|
|
timestep_scale=timestep_scale,
|
|
)
|
|
|
|
# Initialize transformer blocks
|
|
self.noise_refiner = nn.ModuleList(
|
|
[
|
|
OmniGen2TransformerBlock(
|
|
hidden_size,
|
|
num_attention_heads,
|
|
num_kv_heads,
|
|
multiple_of,
|
|
ffn_dim_multiplier,
|
|
norm_eps,
|
|
modulation=True,
|
|
quant_config=quant_config,
|
|
prefix=f"noise_refiner.{i}",
|
|
)
|
|
for i in range(num_refiner_layers)
|
|
]
|
|
)
|
|
|
|
self.ref_image_refiner = nn.ModuleList(
|
|
[
|
|
OmniGen2TransformerBlock(
|
|
hidden_size,
|
|
num_attention_heads,
|
|
num_kv_heads,
|
|
multiple_of,
|
|
ffn_dim_multiplier,
|
|
norm_eps,
|
|
modulation=True,
|
|
quant_config=quant_config,
|
|
prefix=f"ref_image_refiner.{i}",
|
|
)
|
|
for i in range(num_refiner_layers)
|
|
]
|
|
)
|
|
|
|
self.context_refiner = nn.ModuleList(
|
|
[
|
|
OmniGen2TransformerBlock(
|
|
hidden_size,
|
|
num_attention_heads,
|
|
num_kv_heads,
|
|
multiple_of,
|
|
ffn_dim_multiplier,
|
|
norm_eps,
|
|
modulation=False,
|
|
quant_config=quant_config,
|
|
prefix=f"context_refiner.{i}",
|
|
)
|
|
for i in range(num_refiner_layers)
|
|
]
|
|
)
|
|
|
|
# 3. Transformer blocks
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
OmniGen2TransformerBlock(
|
|
hidden_size,
|
|
num_attention_heads,
|
|
num_kv_heads,
|
|
multiple_of,
|
|
ffn_dim_multiplier,
|
|
norm_eps,
|
|
modulation=True,
|
|
quant_config=quant_config,
|
|
prefix=f"layers.{i}",
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
|
|
# 4. Output norm & projection
|
|
self.norm_out = LuminaLayerNormContinuous(
|
|
embedding_dim=hidden_size,
|
|
conditioning_embedding_dim=min(hidden_size, 1024),
|
|
elementwise_affine=False,
|
|
eps=1e-6,
|
|
bias=True,
|
|
out_dim=patch_size * patch_size * self.out_channels,
|
|
)
|
|
|
|
# Add learnable embeddings to distinguish different images
|
|
self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images
|
|
|
|
self.sp_input_boundary = _OmniGen2SPInputBoundary()
|
|
self.sp_output_boundary = _OmniGen2SPOutputBoundary()
|
|
|
|
def img_patch_embed_and_refine(
|
|
self,
|
|
hidden_states,
|
|
ref_image_hidden_states,
|
|
padded_img_mask,
|
|
padded_ref_img_mask,
|
|
noise_rotary_emb,
|
|
ref_img_rotary_emb,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
temb,
|
|
):
|
|
batch_size = len(hidden_states)
|
|
has_ref_tokens = any(ref_img_len > 0 for ref_lens in l_effective_ref_img_len for ref_img_len in ref_lens)
|
|
max_combined_img_len = max(
|
|
[img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len)]
|
|
)
|
|
|
|
hidden_states = self.x_embedder(hidden_states)
|
|
if not has_ref_tokens:
|
|
# FP8 kernels do not support zero-token GEMM on ref_image_patch_embedder; skip that path only.
|
|
# Still run noise_refiner and return the same combined layout as the no-ref case below
|
|
# (batch, max_combined_img_len, hidden) — not raw noise tokens alone.
|
|
for layer in self.noise_refiner:
|
|
hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb)
|
|
combined_img_hidden_states = hidden_states.new_zeros(
|
|
batch_size, max_combined_img_len, self.config.hidden_size
|
|
)
|
|
for i, img_len in enumerate(l_effective_img_len):
|
|
combined_img_hidden_states[i, :img_len] = hidden_states[i, :img_len]
|
|
return combined_img_hidden_states
|
|
|
|
ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states)
|
|
|
|
for i in range(batch_size):
|
|
shift = 0
|
|
for j, ref_img_len in enumerate(l_effective_ref_img_len[i]):
|
|
ref_image_hidden_states[i, shift : shift + ref_img_len, :] = (
|
|
ref_image_hidden_states[i, shift : shift + ref_img_len, :] + self.image_index_embedding[j]
|
|
)
|
|
shift += ref_img_len
|
|
|
|
for layer in self.noise_refiner:
|
|
hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb)
|
|
|
|
flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len))
|
|
num_ref_images = len(flat_l_effective_ref_img_len)
|
|
max_ref_img_len = max(flat_l_effective_ref_img_len)
|
|
|
|
batch_ref_img_mask = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, dtype=torch.bool)
|
|
batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros(
|
|
num_ref_images, max_ref_img_len, self.config.hidden_size
|
|
)
|
|
batch_ref_img_rotary_emb = hidden_states.new_zeros(
|
|
num_ref_images,
|
|
max_ref_img_len,
|
|
ref_img_rotary_emb.shape[-1],
|
|
dtype=ref_img_rotary_emb.dtype,
|
|
)
|
|
batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype)
|
|
|
|
# sequence of ref imgs to batch
|
|
idx = 0
|
|
for i in range(batch_size):
|
|
shift = 0
|
|
for ref_img_len in l_effective_ref_img_len[i]:
|
|
batch_ref_img_mask[idx, :ref_img_len] = True
|
|
batch_ref_image_hidden_states[idx, :ref_img_len] = ref_image_hidden_states[
|
|
i, shift : shift + ref_img_len
|
|
]
|
|
batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[i, shift : shift + ref_img_len]
|
|
batch_temb[idx] = temb[i]
|
|
shift += ref_img_len
|
|
idx += 1
|
|
|
|
# refine ref imgs separately
|
|
for layer in self.ref_image_refiner:
|
|
batch_ref_image_hidden_states = layer(
|
|
batch_ref_image_hidden_states,
|
|
batch_ref_img_mask,
|
|
batch_ref_img_rotary_emb,
|
|
batch_temb,
|
|
)
|
|
|
|
# batch of ref imgs to sequence
|
|
idx = 0
|
|
for i in range(batch_size):
|
|
shift = 0
|
|
for ref_img_len in l_effective_ref_img_len[i]:
|
|
ref_image_hidden_states[i, shift : shift + ref_img_len] = batch_ref_image_hidden_states[
|
|
idx, :ref_img_len
|
|
]
|
|
shift += ref_img_len
|
|
idx += 1
|
|
|
|
combined_img_hidden_states = hidden_states.new_zeros(batch_size, max_combined_img_len, self.config.hidden_size)
|
|
for i, (ref_img_len, img_len) in enumerate(zip(l_effective_ref_img_len, l_effective_img_len)):
|
|
combined_img_hidden_states[i, : sum(ref_img_len)] = ref_image_hidden_states[i, : sum(ref_img_len)]
|
|
combined_img_hidden_states[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = hidden_states[i, :img_len]
|
|
|
|
return combined_img_hidden_states
|
|
|
|
def flat_and_pad_to_seq(self, hidden_states, ref_image_hidden_states):
|
|
batch_size = len(hidden_states)
|
|
p = self.config.patch_size
|
|
device = hidden_states[0].device
|
|
|
|
img_sizes = [(img.size(1), img.size(2)) for img in hidden_states]
|
|
l_effective_img_len = [(H // p) * (W // p) for (H, W) in img_sizes]
|
|
|
|
if ref_image_hidden_states is not None:
|
|
ref_img_sizes = [
|
|
([(img.size(1), img.size(2)) for img in imgs] if imgs is not None else None)
|
|
for imgs in ref_image_hidden_states
|
|
]
|
|
l_effective_ref_img_len = [
|
|
(
|
|
[(ref_img_size[0] // p) * (ref_img_size[1] // p) for ref_img_size in _ref_img_sizes]
|
|
if _ref_img_sizes is not None
|
|
else [0]
|
|
)
|
|
for _ref_img_sizes in ref_img_sizes
|
|
]
|
|
else:
|
|
ref_img_sizes = [None for _ in range(batch_size)]
|
|
l_effective_ref_img_len = [[0] for _ in range(batch_size)]
|
|
|
|
max_ref_img_len = max([sum(ref_img_len) for ref_img_len in l_effective_ref_img_len])
|
|
max_img_len = max(l_effective_img_len)
|
|
|
|
# ref image patch embeddings
|
|
flat_ref_img_hidden_states = []
|
|
for i in range(batch_size):
|
|
if ref_img_sizes[i] is not None:
|
|
imgs = []
|
|
for ref_img in ref_image_hidden_states[i]:
|
|
C, H, W = ref_img.size()
|
|
ref_img = rearrange(ref_img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p)
|
|
imgs.append(ref_img)
|
|
|
|
img = torch.cat(imgs, dim=0)
|
|
flat_ref_img_hidden_states.append(img)
|
|
else:
|
|
flat_ref_img_hidden_states.append(None)
|
|
|
|
# image patch embeddings
|
|
flat_hidden_states = []
|
|
for i in range(batch_size):
|
|
img = hidden_states[i]
|
|
C, H, W = img.size()
|
|
|
|
img = rearrange(img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p)
|
|
flat_hidden_states.append(img)
|
|
|
|
padded_ref_img_hidden_states = torch.zeros(
|
|
batch_size,
|
|
max_ref_img_len,
|
|
flat_hidden_states[0].shape[-1],
|
|
device=device,
|
|
dtype=flat_hidden_states[0].dtype,
|
|
)
|
|
padded_ref_img_mask = torch.zeros(batch_size, max_ref_img_len, dtype=torch.bool, device=device)
|
|
for i in range(batch_size):
|
|
if ref_img_sizes[i] is not None:
|
|
padded_ref_img_hidden_states[i, : sum(l_effective_ref_img_len[i])] = flat_ref_img_hidden_states[i]
|
|
padded_ref_img_mask[i, : sum(l_effective_ref_img_len[i])] = True
|
|
|
|
padded_hidden_states = torch.zeros(
|
|
batch_size,
|
|
max_img_len,
|
|
flat_hidden_states[0].shape[-1],
|
|
device=device,
|
|
dtype=flat_hidden_states[0].dtype,
|
|
)
|
|
padded_img_mask = torch.zeros(batch_size, max_img_len, dtype=torch.bool, device=device)
|
|
for i in range(batch_size):
|
|
padded_hidden_states[i, : l_effective_img_len[i]] = flat_hidden_states[i]
|
|
padded_img_mask[i, : l_effective_img_len[i]] = True
|
|
|
|
return (
|
|
padded_hidden_states,
|
|
padded_ref_img_hidden_states,
|
|
padded_img_mask,
|
|
padded_ref_img_mask,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
ref_img_sizes,
|
|
img_sizes,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor | list[torch.Tensor],
|
|
timestep: torch.Tensor,
|
|
text_hidden_states: torch.Tensor,
|
|
freqs_cis: torch.Tensor,
|
|
text_attention_mask: torch.Tensor,
|
|
ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
|
|
return_dict: bool = False,
|
|
) -> torch.Tensor | Transformer2DModelOutput:
|
|
# 1. Condition, positional & patch embedding
|
|
batch_size = len(hidden_states)
|
|
is_hidden_states_tensor = isinstance(hidden_states, torch.Tensor)
|
|
|
|
if is_hidden_states_tensor:
|
|
assert hidden_states.ndim == 4
|
|
hidden_states = [_hidden_states for _hidden_states in hidden_states]
|
|
|
|
device = hidden_states[0].device
|
|
|
|
temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states[0].dtype)
|
|
|
|
(
|
|
hidden_states,
|
|
ref_image_hidden_states,
|
|
img_mask,
|
|
ref_img_mask,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
ref_img_sizes,
|
|
img_sizes,
|
|
) = self.flat_and_pad_to_seq(hidden_states, ref_image_hidden_states)
|
|
|
|
(
|
|
context_rotary_emb,
|
|
ref_img_rotary_emb,
|
|
noise_rotary_emb,
|
|
rotary_emb,
|
|
encoder_seq_lengths,
|
|
seq_lengths,
|
|
) = self.rope_embedder(
|
|
freqs_cis,
|
|
text_attention_mask,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
ref_img_sizes,
|
|
img_sizes,
|
|
device,
|
|
)
|
|
|
|
# 2. Context refinement
|
|
for layer in self.context_refiner:
|
|
text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb)
|
|
|
|
combined_img_hidden_states = self.img_patch_embed_and_refine(
|
|
hidden_states,
|
|
ref_image_hidden_states,
|
|
img_mask,
|
|
ref_img_mask,
|
|
noise_rotary_emb,
|
|
ref_img_rotary_emb,
|
|
l_effective_ref_img_len,
|
|
l_effective_img_len,
|
|
temb,
|
|
)
|
|
|
|
# 3. Joint Transformer blocks
|
|
max_seq_len = max(seq_lengths)
|
|
|
|
attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
|
|
joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size)
|
|
for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
|
|
attention_mask[i, :seq_len] = True
|
|
joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
|
|
joint_hidden_states[i, encoder_seq_len:seq_len] = combined_img_hidden_states[i, : seq_len - encoder_seq_len]
|
|
|
|
hidden_states = joint_hidden_states
|
|
|
|
hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb)
|
|
|
|
for layer in self.layers:
|
|
hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
|
|
|
|
hidden_states = self.sp_output_boundary(hidden_states)
|
|
|
|
# 4. Output norm & projection
|
|
hidden_states = self.norm_out(hidden_states, temb)
|
|
|
|
p = self.config.patch_size
|
|
output = []
|
|
for i, (img_size, img_len, seq_len) in enumerate(zip(img_sizes, l_effective_img_len, seq_lengths)):
|
|
height, width = img_size
|
|
output.append(
|
|
rearrange(
|
|
hidden_states[i][seq_len - img_len : seq_len],
|
|
"(h w) (p1 p2 c) -> c (h p1) (w p2)",
|
|
h=height // p,
|
|
w=width // p,
|
|
p1=p,
|
|
p2=p,
|
|
)
|
|
)
|
|
if is_hidden_states_tensor:
|
|
output = torch.stack(output, dim=0)
|
|
|
|
if not return_dict:
|
|
return output
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
stacked_params_mapping = [
|
|
# (param_name, shard_name, shard_id)
|
|
# self-attn
|
|
(".to_qkv", ".to_q", "q"),
|
|
(".to_qkv", ".to_k", "k"),
|
|
(".to_qkv", ".to_v", "v"),
|
|
# feed-forward
|
|
(".feed_forward.gate_up_proj", ".feed_forward.linear_1", 0),
|
|
(".feed_forward.gate_up_proj", ".feed_forward.linear_3", 1),
|
|
]
|
|
|
|
params_dict = dict(self.named_parameters())
|
|
|
|
loaded_params = set[str]()
|
|
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:
|
|
# feed_forward.linear_2 -> feed_forward.down_proj rename
|
|
if ".feed_forward.linear_2." in name:
|
|
name = name.replace(".feed_forward.linear_2.", ".feed_forward.down_proj.")
|
|
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
|