721 lines
28 KiB
Python
721 lines
28 KiB
Python
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import functools
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.models.attention import FeedForward
|
|
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
from diffusers.models.normalization import AdaLayerNormContinuous
|
|
|
|
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
|
|
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
|
|
from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
|
|
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
|
|
from sglang.multimodal_gen.runtime.layers.triton_ops import (
|
|
apply_rotary_embedding,
|
|
fuse_scale_shift_kernel,
|
|
)
|
|
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
|
|
from sglang.multimodal_gen.runtime.models.dits.utils import (
|
|
delete_projection_layers,
|
|
fuse_linear_projections,
|
|
)
|
|
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
|
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
|
|
|
logger = init_logger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
def _get_projections(
|
|
attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
|
|
):
|
|
img_query, _ = attn.to_q(hidden_states)
|
|
img_key, _ = attn.to_k(hidden_states)
|
|
img_value, _ = attn.to_v(hidden_states)
|
|
|
|
txt_query = txt_key = txt_value = None
|
|
if encoder_hidden_states is not None and hasattr(attn, "add_q_proj"):
|
|
txt_query, _ = attn.add_q_proj(encoder_hidden_states)
|
|
txt_key, _ = attn.add_k_proj(encoder_hidden_states)
|
|
txt_value, _ = attn.add_v_proj(encoder_hidden_states)
|
|
|
|
return img_query, img_key, img_value, txt_query, txt_key, txt_value
|
|
|
|
|
|
def _get_fused_projections(
|
|
attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
|
|
):
|
|
img_qkv, _ = attn.to_qkv(hidden_states)
|
|
img_query, img_key, img_value = img_qkv.chunk(3, dim=-1)
|
|
|
|
txt_query = txt_key = txt_value = None
|
|
if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
|
|
txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
|
|
txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1)
|
|
|
|
return img_query, img_key, img_value, txt_query, txt_key, txt_value
|
|
|
|
|
|
def _get_qkv_projections(
|
|
attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
|
|
):
|
|
if attn.fused_projections:
|
|
return _get_fused_projections(attn, hidden_states, encoder_hidden_states)
|
|
return _get_projections(attn, hidden_states, encoder_hidden_states)
|
|
|
|
|
|
class QwenTimestepProjEmbeddings(nn.Module):
|
|
def __init__(self, embedding_dim):
|
|
super().__init__()
|
|
|
|
self.time_proj = Timesteps(
|
|
num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000
|
|
)
|
|
self.timestep_embedder = TimestepEmbedding(
|
|
in_channels=256, time_embed_dim=embedding_dim
|
|
)
|
|
|
|
def forward(self, timestep, hidden_states):
|
|
timesteps_proj = self.time_proj(timestep)
|
|
timesteps_emb = self.timestep_embedder(
|
|
timesteps_proj.to(dtype=hidden_states.dtype)
|
|
) # (N, D)
|
|
|
|
conditioning = timesteps_emb
|
|
|
|
return conditioning
|
|
|
|
|
|
class QwenEmbedRope(nn.Module):
|
|
def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
|
|
super().__init__()
|
|
self.theta = theta
|
|
self.axes_dim = axes_dim
|
|
pos_index = torch.arange(4096)
|
|
neg_index = torch.arange(4096).flip(0) * -1 - 1
|
|
self.pos_freqs = torch.cat(
|
|
[
|
|
self.rope_params(pos_index, self.axes_dim[0], self.theta),
|
|
self.rope_params(pos_index, self.axes_dim[1], self.theta),
|
|
self.rope_params(pos_index, self.axes_dim[2], self.theta),
|
|
],
|
|
dim=1,
|
|
)
|
|
self.neg_freqs = torch.cat(
|
|
[
|
|
self.rope_params(neg_index, self.axes_dim[0], self.theta),
|
|
self.rope_params(neg_index, self.axes_dim[1], self.theta),
|
|
self.rope_params(neg_index, self.axes_dim[2], self.theta),
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
# self.rope = NDRotaryEmbedding(
|
|
# rope_dim_list=axes_dim,
|
|
# rope_theta=theta,
|
|
# use_real=False,
|
|
# repeat_interleave_real=False,
|
|
# dtype=torch.float32 if current_platform.is_mps() else torch.float64,
|
|
# )
|
|
|
|
# DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
|
|
self.scale_rope = scale_rope
|
|
|
|
def rope_params(self, index, dim, theta=10000):
|
|
"""
|
|
Args:
|
|
index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
|
|
"""
|
|
device = index.device
|
|
assert dim % 2 == 0
|
|
freqs = torch.outer(
|
|
index,
|
|
(
|
|
1.0
|
|
/ torch.pow(
|
|
theta,
|
|
torch.arange(0, dim, 2, device=device).to(torch.float32).div(dim),
|
|
)
|
|
).to(device=device),
|
|
)
|
|
freqs = torch.polar(torch.ones_like(freqs), freqs)
|
|
return freqs
|
|
|
|
def forward(
|
|
self,
|
|
video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
|
|
txt_seq_lens: List[int],
|
|
device: torch.device,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
|
|
A list of 3 integers [frame, height, width] representing the shape of the video.
|
|
txt_seq_lens (`List[int]`):
|
|
A list of integers of length batch_size representing the length of each text prompt.
|
|
device: (`torch.device`):
|
|
The device on which to perform the RoPE computation.
|
|
"""
|
|
# When models are initialized under a "meta" device context (e.g. init_empty_weights),
|
|
# tensors created during __init__ become meta tensors. Calling .to(...) on a meta tensor
|
|
# raises "Cannot copy out of meta tensor". Rebuild the frequencies on the target device
|
|
# in that case; otherwise move them if just on a different device.
|
|
if getattr(self.pos_freqs, "device", torch.device("meta")).type == "meta":
|
|
pos_index = torch.arange(4096, device=device)
|
|
neg_index = torch.arange(4096, device=device).flip(0) * -1 - 1
|
|
self.pos_freqs = torch.cat(
|
|
[
|
|
self.rope_params(pos_index, self.axes_dim[0], self.theta),
|
|
self.rope_params(pos_index, self.axes_dim[1], self.theta),
|
|
self.rope_params(pos_index, self.axes_dim[2], self.theta),
|
|
],
|
|
dim=1,
|
|
).to(device=device)
|
|
self.neg_freqs = torch.cat(
|
|
[
|
|
self.rope_params(neg_index, self.axes_dim[0], self.theta),
|
|
self.rope_params(neg_index, self.axes_dim[1], self.theta),
|
|
self.rope_params(neg_index, self.axes_dim[2], self.theta),
|
|
],
|
|
dim=1,
|
|
).to(device=device)
|
|
elif self.pos_freqs.device != device:
|
|
self.pos_freqs = self.pos_freqs.to(device)
|
|
self.neg_freqs = self.neg_freqs.to(device)
|
|
|
|
if isinstance(video_fhw, list):
|
|
video_fhw = video_fhw[0]
|
|
if not isinstance(video_fhw, list):
|
|
video_fhw = [video_fhw]
|
|
|
|
vid_freqs = []
|
|
max_vid_index = 0
|
|
for idx, fhw in enumerate(video_fhw):
|
|
frame, height, width = fhw
|
|
# RoPE frequencies are cached via a lru_cache decorator on _compute_video_freqs
|
|
video_freq = self._compute_video_freqs(frame, height, width, idx)
|
|
video_freq = video_freq.to(device)
|
|
vid_freqs.append(video_freq)
|
|
|
|
if self.scale_rope:
|
|
max_vid_index = max(height // 2, width // 2, max_vid_index)
|
|
else:
|
|
max_vid_index = max(height, width, max_vid_index)
|
|
|
|
max_len = max(txt_seq_lens)
|
|
txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
|
|
vid_freqs = torch.cat(vid_freqs, dim=0).to(device=device)
|
|
return vid_freqs, txt_freqs
|
|
|
|
@functools.lru_cache(maxsize=128)
|
|
def _compute_video_freqs(
|
|
self, frame: int, height: int, width: int, idx: int = 0
|
|
) -> torch.Tensor:
|
|
seq_lens = frame * height * width
|
|
freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
|
|
freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
|
|
|
|
freqs_frame = (
|
|
freqs_pos[0][idx : idx + frame]
|
|
.view(frame, 1, 1, -1)
|
|
.expand(frame, height, width, -1)
|
|
)
|
|
if self.scale_rope:
|
|
freqs_height = torch.cat(
|
|
[freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
|
|
dim=0,
|
|
)
|
|
freqs_height = freqs_height.view(1, height, 1, -1).expand(
|
|
frame, height, width, -1
|
|
)
|
|
freqs_width = torch.cat(
|
|
[freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
|
|
dim=0,
|
|
)
|
|
freqs_width = freqs_width.view(1, 1, width, -1).expand(
|
|
frame, height, width, -1
|
|
)
|
|
else:
|
|
freqs_height = (
|
|
freqs_pos[1][:height]
|
|
.view(1, height, 1, -1)
|
|
.expand(frame, height, width, -1)
|
|
)
|
|
freqs_width = (
|
|
freqs_pos[2][:width]
|
|
.view(1, 1, width, -1)
|
|
.expand(frame, height, width, -1)
|
|
)
|
|
|
|
freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(
|
|
seq_lens, -1
|
|
)
|
|
return freqs.clone().contiguous()
|
|
|
|
|
|
class QwenImageCrossAttention(nn.Module):
|
|
_supports_qkv_fusion = True
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int, # query_dim
|
|
num_heads: int,
|
|
head_dim: int,
|
|
window_size=(-1, -1),
|
|
added_kv_proj_dim: int = None,
|
|
out_bias: bool = True,
|
|
qk_norm=True, # rmsnorm
|
|
eps=1e-6,
|
|
pre_only=False,
|
|
context_pre_only: bool = False,
|
|
parallel_attention=False,
|
|
out_dim: int = None,
|
|
) -> None:
|
|
assert dim % num_heads == 0
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.window_size = window_size
|
|
self.qk_norm = qk_norm
|
|
self.eps = eps
|
|
self.parallel_attention = parallel_attention
|
|
|
|
# layers
|
|
self.to_q = ReplicatedLinear(dim, dim)
|
|
self.to_k = ReplicatedLinear(dim, dim)
|
|
self.to_v = ReplicatedLinear(dim, dim)
|
|
if self.qk_norm:
|
|
self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
|
|
self.inner_kv_dim = self.inner_dim
|
|
if added_kv_proj_dim is not None:
|
|
self.add_k_proj = ReplicatedLinear(
|
|
added_kv_proj_dim, self.inner_kv_dim, bias=True
|
|
)
|
|
self.add_v_proj = ReplicatedLinear(
|
|
added_kv_proj_dim, self.inner_kv_dim, bias=True
|
|
)
|
|
if context_pre_only is not None:
|
|
self.add_q_proj = ReplicatedLinear(
|
|
added_kv_proj_dim, self.inner_dim, bias=True
|
|
)
|
|
|
|
if context_pre_only is not None and not context_pre_only:
|
|
self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
|
|
else:
|
|
self.to_add_out = None
|
|
|
|
if not pre_only:
|
|
self.to_out = nn.ModuleList([])
|
|
self.to_out.append(
|
|
ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
|
|
)
|
|
else:
|
|
self.to_out = None
|
|
|
|
self.norm_added_q = RMSNorm(head_dim, eps=eps)
|
|
self.norm_added_k = RMSNorm(head_dim, eps=eps)
|
|
|
|
# Scaled dot product attention
|
|
self.attn = USPAttention(
|
|
num_heads=num_heads,
|
|
head_size=self.head_dim,
|
|
dropout_rate=0,
|
|
softmax_scale=None,
|
|
causal=False,
|
|
supported_attention_backends={
|
|
AttentionBackendEnum.FA,
|
|
AttentionBackendEnum.TORCH_SDPA,
|
|
AttentionBackendEnum.SAGE_ATTN,
|
|
},
|
|
)
|
|
|
|
self.fused_projections = False
|
|
self.added_kv_proj_dim_val = added_kv_proj_dim
|
|
|
|
@torch.no_grad()
|
|
def fuse_projections(self):
|
|
if self.fused_projections:
|
|
return
|
|
|
|
self.to_qkv = fuse_linear_projections(
|
|
self.to_q, self.to_k, self.to_v, use_bias=False, linear_cls=ReplicatedLinear
|
|
)
|
|
delete_projection_layers(self, ["to_q", "to_k", "to_v"])
|
|
|
|
if self.added_kv_proj_dim_val is not None and hasattr(self, "add_q_proj"):
|
|
self.to_added_qkv = fuse_linear_projections(
|
|
self.add_q_proj,
|
|
self.add_k_proj,
|
|
self.add_v_proj,
|
|
use_bias=True,
|
|
linear_cls=ReplicatedLinear,
|
|
)
|
|
delete_projection_layers(self, ["add_q_proj", "add_k_proj", "add_v_proj"])
|
|
|
|
self.fused_projections = True
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
|
**cross_attention_kwargs,
|
|
):
|
|
seq_len_txt = encoder_hidden_states.shape[1]
|
|
|
|
img_query, img_key, img_value, txt_query, txt_key, txt_value = (
|
|
_get_qkv_projections(self, hidden_states, encoder_hidden_states)
|
|
)
|
|
|
|
# Reshape for multi-head attention
|
|
img_query = img_query.unflatten(-1, (self.num_heads, -1))
|
|
img_key = img_key.unflatten(-1, (self.num_heads, -1))
|
|
img_value = img_value.unflatten(-1, (self.num_heads, -1))
|
|
|
|
txt_query = txt_query.unflatten(-1, (self.num_heads, -1))
|
|
txt_key = txt_key.unflatten(-1, (self.num_heads, -1))
|
|
txt_value = txt_value.unflatten(-1, (self.num_heads, -1))
|
|
|
|
# Apply QK normalization
|
|
if self.norm_q is not None:
|
|
img_query = self.norm_q(img_query)
|
|
if self.norm_k is not None:
|
|
img_key = self.norm_k(img_key)
|
|
if self.norm_added_q is not None:
|
|
txt_query = self.norm_added_q(txt_query)
|
|
if self.norm_added_k is not None:
|
|
txt_key = self.norm_added_k(txt_key)
|
|
|
|
# Apply RoPE
|
|
if image_rotary_emb is not None:
|
|
(img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
|
|
img_query = apply_rotary_embedding(
|
|
img_query, img_cos, img_sin, interleaved=True
|
|
)
|
|
img_key = apply_rotary_embedding(
|
|
img_key, img_cos, img_sin, interleaved=True
|
|
)
|
|
txt_query = apply_rotary_embedding(
|
|
txt_query, txt_cos, txt_sin, interleaved=True
|
|
)
|
|
txt_key = apply_rotary_embedding(
|
|
txt_key, txt_cos, txt_sin, interleaved=True
|
|
)
|
|
|
|
# Concatenate for joint attention
|
|
# Order: [text, image]
|
|
joint_query = torch.cat([txt_query, img_query], dim=1)
|
|
joint_key = torch.cat([txt_key, img_key], dim=1)
|
|
joint_value = torch.cat([txt_value, img_value], dim=1)
|
|
|
|
# Compute joint attention
|
|
joint_hidden_states = self.attn(
|
|
joint_query,
|
|
joint_key,
|
|
joint_value,
|
|
)
|
|
|
|
# Reshape back
|
|
joint_hidden_states = joint_hidden_states.flatten(2, 3)
|
|
joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
|
|
|
|
# Split attention outputs back
|
|
txt_attn_output = joint_hidden_states[:, :seq_len_txt, :] # Text part
|
|
img_attn_output = joint_hidden_states[:, seq_len_txt:, :] # Image part
|
|
|
|
# Apply output projections
|
|
img_attn_output, _ = self.to_out[0](img_attn_output)
|
|
if len(self.to_out) > 1:
|
|
(img_attn_output,) = self.to_out[1](img_attn_output) # dropout
|
|
|
|
txt_attn_output, _ = self.to_add_out(txt_attn_output)
|
|
|
|
return img_attn_output, txt_attn_output
|
|
|
|
|
|
class QwenImageTransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_attention_heads: int,
|
|
attention_head_dim: int,
|
|
qk_norm: str = "rms_norm",
|
|
eps: float = 1e-6,
|
|
):
|
|
super().__init__()
|
|
|
|
self.dim = dim
|
|
self.num_attention_heads = num_attention_heads
|
|
self.attention_head_dim = attention_head_dim
|
|
|
|
# Image processing modules
|
|
self.img_mod = nn.Sequential(
|
|
nn.SiLU(),
|
|
nn.Linear(
|
|
dim, 6 * dim, bias=True
|
|
), # For scale, shift, gate for norm1 and norm2
|
|
)
|
|
self.img_norm1 = LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
|
|
self.attn = QwenImageCrossAttention(
|
|
dim=dim,
|
|
num_heads=num_attention_heads,
|
|
added_kv_proj_dim=dim,
|
|
context_pre_only=False,
|
|
head_dim=attention_head_dim,
|
|
)
|
|
self.img_norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False)
|
|
self.img_mlp = FeedForward(
|
|
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
|
|
)
|
|
|
|
# Text processing modules
|
|
self.txt_mod = nn.Sequential(
|
|
nn.SiLU(),
|
|
nn.Linear(
|
|
dim, 6 * dim, bias=True
|
|
), # For scale, shift, gate for norm1 and norm2
|
|
)
|
|
self.txt_norm1 = LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
# Text doesn't need separate attention - it's handled by img_attn joint computation
|
|
self.txt_norm2 = LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
self.txt_mlp = FeedForward(
|
|
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
|
|
)
|
|
|
|
def _modulate(self, x, mod_params):
|
|
"""Apply modulation to input tensor"""
|
|
shift, scale, gate = mod_params.chunk(3, dim=-1)
|
|
return fuse_scale_shift_kernel(x, scale, shift), gate.unsqueeze(1)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
encoder_hidden_states_mask: torch.Tensor,
|
|
temb: torch.Tensor,
|
|
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
|
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
# Get modulation parameters for both streams
|
|
img_mod_params = self.img_mod(temb) # [B, 6*dim]
|
|
txt_mod_params = self.txt_mod(temb) # [B, 6*dim]
|
|
|
|
# Split modulation parameters for norm1 and norm2
|
|
img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
|
|
txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
|
|
|
|
# Process image stream - norm1 + modulation
|
|
|
|
img_normed = self.img_norm1(hidden_states)
|
|
|
|
img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
|
|
|
|
# Process text stream - norm1 + modulation
|
|
txt_normed = self.txt_norm1(encoder_hidden_states)
|
|
txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
|
|
|
|
# Use QwenAttnProcessor2_0 for joint attention computation
|
|
# This directly implements the DoubleStreamLayerMegatron logic:
|
|
# 1. Computes QKV for both streams
|
|
# 2. Applies QK normalization and RoPE
|
|
# 3. Concatenates and runs joint attention
|
|
# 4. Splits results back to separate streams
|
|
joint_attention_kwargs = joint_attention_kwargs or {}
|
|
attn_output = self.attn(
|
|
hidden_states=img_modulated, # Image stream (will be processed as "sample")
|
|
encoder_hidden_states=txt_modulated, # Text stream (will be processed as "context")
|
|
encoder_hidden_states_mask=encoder_hidden_states_mask,
|
|
image_rotary_emb=image_rotary_emb,
|
|
**joint_attention_kwargs,
|
|
)
|
|
|
|
# QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
|
|
img_attn_output, txt_attn_output = attn_output
|
|
|
|
# Apply attention gates and add residual (like in Megatron)
|
|
hidden_states = hidden_states + img_gate1 * img_attn_output
|
|
|
|
encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
|
|
|
|
# Process image stream - norm2 + MLP
|
|
img_normed2 = self.img_norm2(hidden_states)
|
|
img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
|
|
img_mlp_output = self.img_mlp(img_modulated2)
|
|
hidden_states = hidden_states + img_gate2 * img_mlp_output
|
|
|
|
# Process text stream - norm2 + MLP
|
|
txt_normed2 = self.txt_norm2(encoder_hidden_states)
|
|
txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
|
|
txt_mlp_output = self.txt_mlp(txt_modulated2)
|
|
encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
|
|
|
|
# Clip to prevent overflow for fp16
|
|
if encoder_hidden_states.dtype == torch.float16:
|
|
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
|
|
if hidden_states.dtype == torch.float16:
|
|
hidden_states = hidden_states.clip(-65504, 65504)
|
|
|
|
return encoder_hidden_states, hidden_states
|
|
|
|
|
|
class QwenImageTransformer2DModel(CachableDiT):
|
|
"""
|
|
The Transformer model introduced in Qwen.
|
|
|
|
"""
|
|
|
|
_supports_gradient_checkpointing = True
|
|
_no_split_modules = ["QwenImageTransformerBlock"]
|
|
_skip_layerwise_casting_patterns = ["pos_embed", "norm"]
|
|
_repeated_blocks = ["QwenImageTransformerBlock"]
|
|
|
|
def __init__(
|
|
self,
|
|
config: QwenImageDitConfig,
|
|
hf_config: dict[str, Any],
|
|
):
|
|
super().__init__(config=config, hf_config=hf_config)
|
|
patch_size = config.arch_config.patch_size
|
|
in_channels = config.arch_config.in_channels
|
|
out_channels = config.arch_config.out_channels
|
|
num_layers = config.arch_config.num_layers
|
|
attention_head_dim = config.arch_config.attention_head_dim
|
|
num_attention_heads = config.arch_config.num_attention_heads
|
|
joint_attention_dim = config.arch_config.joint_attention_dim
|
|
axes_dims_rope = config.arch_config.axes_dims_rope
|
|
self.out_channels = out_channels or in_channels
|
|
self.inner_dim = num_attention_heads * attention_head_dim
|
|
|
|
self.rotary_emb = QwenEmbedRope(
|
|
theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True
|
|
)
|
|
|
|
self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
|
|
|
|
self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
|
|
|
|
self.img_in = nn.Linear(in_channels, self.inner_dim)
|
|
self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
|
|
|
|
self.transformer_blocks = nn.ModuleList(
|
|
[
|
|
QwenImageTransformerBlock(
|
|
dim=self.inner_dim,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=attention_head_dim,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
|
|
self.norm_out = AdaLayerNormContinuous(
|
|
self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
|
|
)
|
|
self.proj_out = nn.Linear(
|
|
self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
|
|
)
|
|
|
|
def fuse_qkv_projections(self):
|
|
for block in self.transformer_blocks:
|
|
if hasattr(block.attn, "fuse_projections") and getattr(
|
|
block.attn, "_supports_qkv_fusion", True
|
|
):
|
|
block.attn.fuse_projections()
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor = None,
|
|
encoder_hidden_states_mask: torch.Tensor = None,
|
|
timestep: torch.LongTensor = None,
|
|
txt_seq_lens: Optional[List[int]] = None,
|
|
freqs_cis: tuple[torch.Tensor, torch.Tensor] = None,
|
|
guidance: torch.Tensor = None, # TODO: this should probably be removed
|
|
attention_kwargs: Optional[Dict[str, Any]] = None,
|
|
controlnet_block_samples=None,
|
|
return_dict: bool = True,
|
|
) -> Union[torch.Tensor, Transformer2DModelOutput]:
|
|
"""
|
|
The [`QwenTransformer2DModel`] forward method.
|
|
|
|
Args:
|
|
hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
|
|
Input `hidden_states`.
|
|
encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
|
|
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
|
|
encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
|
|
Mask of the input conditions.
|
|
timestep ( `torch.LongTensor`):
|
|
Used to indicate denoising step.
|
|
attention_kwargs (`dict`, *optional*):
|
|
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
|
`self.processor` in
|
|
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
|
|
tuple.
|
|
|
|
Returns:
|
|
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
|
`tuple` where the first element is the sample tensor.
|
|
"""
|
|
if (
|
|
attention_kwargs is not None
|
|
and attention_kwargs.get("scale", None) is not None
|
|
):
|
|
logger.warning(
|
|
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
|
|
)
|
|
|
|
if isinstance(encoder_hidden_states, list):
|
|
encoder_hidden_states = encoder_hidden_states[0]
|
|
|
|
hidden_states = self.img_in(hidden_states)
|
|
|
|
timestep = (timestep / 1000).to(hidden_states.dtype)
|
|
encoder_hidden_states = self.txt_norm(encoder_hidden_states)
|
|
encoder_hidden_states = self.txt_in(encoder_hidden_states)
|
|
|
|
temb = self.time_text_embed(timestep, hidden_states)
|
|
|
|
image_rotary_emb = freqs_cis
|
|
for index_block, block in enumerate(self.transformer_blocks):
|
|
encoder_hidden_states, hidden_states = block(
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
encoder_hidden_states_mask=encoder_hidden_states_mask,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
joint_attention_kwargs=attention_kwargs,
|
|
)
|
|
|
|
# controlnet residual
|
|
if controlnet_block_samples is not None:
|
|
interval_control = len(self.transformer_blocks) / len(
|
|
controlnet_block_samples
|
|
)
|
|
interval_control = int(np.ceil(interval_control))
|
|
hidden_states = (
|
|
hidden_states
|
|
+ controlnet_block_samples[index_block // interval_control]
|
|
)
|
|
|
|
# Use only the image part (hidden_states) from the dual-stream blocks
|
|
hidden_states = self.norm_out(hidden_states, temb)
|
|
|
|
output = self.proj_out(hidden_states)
|
|
return output
|
|
|
|
|
|
EntryClass = QwenImageTransformer2DModel
|