chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] improve: improve qwen-image-edit performance to align with LightX2V (#15812)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2025-12-26 22:25:10 +08:00
committed by GitHub
parent 8dc6f0fc4d
commit 51dbdb2202
5 changed files with 425 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/sglang/multimodal_gen/configs/models/dits/qwenimage.py
View File

@@ -22,27 +22,10 @@ class QwenImageArchConfig(DiTArchConfig):
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
zero_cond_t: bool = False
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
(".to_qkv", ".to_q", "q"),
(".to_qkv", ".to_k", "k"),
(".to_qkv", ".to_v", "v"),
(".to_added_qkv", ".add_q_proj", "q"),
(".to_added_qkv", ".add_k_proj", "k"),
(".to_added_qkv", ".add_v_proj", "v"),
]
)
stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
param_names_mapping: dict = field(
default_factory=lambda: {
# QKV fusion mappings
r"(.*)\.to_q\.(weight|bias)$": (r"\1.to_qkv.\2", 0, 3),
r"(.*)\.to_k\.(weight|bias)$": (r"\1.to_qkv.\2", 1, 3),
r"(.*)\.to_v\.(weight|bias)$": (r"\1.to_qkv.\2", 2, 3),
r"(.*)\.add_q_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 0, 3),
r"(.*)\.add_k_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 1, 3),
r"(.*)\.add_v_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 2, 3),
# LoRA mappings
r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
}

64
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
View File

@@ -156,16 +156,14 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
# img_shapes: for global entire image
img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device)
img_cos, img_sin = (
img_freqs.real.to(dtype=dtype),
img_freqs.imag.to(dtype=dtype),
)
txt_cos, txt_sin = (
txt_freqs.real.to(dtype=dtype),
txt_freqs.imag.to(dtype=dtype),
)
img_cos_half = img_freqs.real.to(dtype=torch.float32).contiguous()
img_sin_half = img_freqs.imag.to(dtype=torch.float32).contiguous()
txt_cos_half = txt_freqs.real.to(dtype=torch.float32).contiguous()
txt_sin_half = txt_freqs.imag.to(dtype=torch.float32).contiguous()
return (img_cos, img_sin), (txt_cos, txt_sin)
img_cos_sin_cache = torch.cat([img_cos_half, img_sin_half], dim=-1)
txt_cos_sin_cache = torch.cat([txt_cos_half, txt_sin_half], dim=-1)
return img_cos_sin_cache, txt_cos_sin_cache
def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype):
batch_size = prompt_embeds[0].shape[0]
@@ -184,15 +182,15 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
(img_cos, img_sin), (txt_cos, txt_sin) = self.get_freqs_cis(
freqs_cis = self.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
img_cos = shard_rotary_emb_for_sp(img_cos)
img_sin = shard_rotary_emb_for_sp(img_sin)
img_cache, txt_cache = freqs_cis
img_cache = shard_rotary_emb_for_sp(img_cache)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
"freqs_cis": (img_cache, txt_cache),
}
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
@@ -252,7 +250,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
],
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
(img_cos, img_sin), (txt_cos, txt_sin) = QwenImagePipelineConfig.get_freqs_cis(
freqs_cis = QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
@@ -261,20 +259,14 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
# concat back the img_cos for input image (since it is not sp-shared later)
img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
img_cache, txt_cache = freqs_cis
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
img_cache = torch.cat(
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
"freqs_cis": (img_cache, txt_cache),
}
def preprocess_condition_image(
@@ -441,16 +433,28 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
(img_cos, img_sin), (txt_cos, txt_sin) = (
QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
# perform sp shard on noisy image tokens
noisy_img_seq_len = (
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
if isinstance(freqs_cis[0], torch.Tensor) and freqs_cis[0].dim() == 2:
img_cache, txt_cache = freqs_cis
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
img_cache = torch.cat(
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": (img_cache, txt_cache),
"img_shapes": img_shapes,
}
(img_cos, img_sin), (txt_cos, txt_sin) = freqs_cis
noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])

96
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
View File

@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass
from functools import lru_cache
from typing import Any
import torch
@@ -18,6 +19,14 @@ try:
except ImportError as e:
raise e
try:
from flash_attn_interface import (
flash_attn_varlen_func as flash_attn_varlen_func_upstream,
)
except Exception:
flash_attn_varlen_func_upstream = None
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
AttentionBackend,
AttentionImpl,
@@ -31,6 +40,51 @@ logger = init_logger(__name__)
fa_ver = 3
@lru_cache(maxsize=128)
def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor:
return torch.arange(
0,
(bsz + 1) * seqlen,
step=seqlen,
device=torch.device("cuda", device_index),
dtype=torch.int32,
)
@lru_cache(maxsize=256)
def _should_use_upstream_flash_attention(
upstream_available: bool,
upstream_heads_ok: bool,
q_shape: tuple[int, ...],
k_shape: tuple[int, ...],
v_shape: tuple[int, ...],
) -> bool:
if not upstream_available or not upstream_heads_ok:
return False
if len(q_shape) != 4 or len(k_shape) != 4 or len(v_shape) != 4:
return False
bsz, seqlen, nheads_q, d = q_shape
bsz_k, seqlen_k, nheads_k, d_k = k_shape
bsz_v, seqlen_v, nheads_v, d_v = v_shape
if (
bsz != bsz_k
or bsz != bsz_v
or seqlen != seqlen_k
or seqlen != seqlen_v
or d != d_k
or d != d_v
):
return False
if nheads_k != nheads_v:
return False
if nheads_k == 0 or (nheads_q % nheads_k) != 0:
return False
return True
def set_fa_ver(ver: int):
global fa_ver
fa_ver = ver
@@ -102,9 +156,19 @@ class FlashAttentionImpl(AttentionImpl):
prefix: str = "",
**extra_impl_args,
) -> None:
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.head_size = head_size
self.causal = causal
self.softmax_scale = softmax_scale
self.attention_metadata = FlashAttentionMetadata()
if self.num_kv_heads is None:
self._upstream_heads_ok = True
else:
# For gqa, the num_heads must be a multiple of num_kv_heads
self._upstream_heads_ok = (
self.num_kv_heads > 0 and (self.num_heads % self.num_kv_heads) == 0
)
def forward(
self,
@@ -124,6 +188,38 @@ class FlashAttentionImpl(AttentionImpl):
else:
max_seqlen_q = query.shape[1]
max_seqlen_k = key.shape[1]
q_shape = tuple(query.shape)
k_shape = tuple(key.shape)
v_shape = tuple(value.shape)
use_upstream = _should_use_upstream_flash_attention(
flash_attn_varlen_func_upstream is not None,
self._upstream_heads_ok,
q_shape,
k_shape,
v_shape,
)
if use_upstream:
bsz, seqlen, nheads_q, d = q_shape
bsz_k, seqlen_k, nheads_k, d_k = k_shape
bsz_v, seqlen_v, nheads_v, d_v = v_shape
q_ = query.contiguous().reshape(bsz * seqlen, nheads_q, d)
k_ = key.contiguous().reshape(bsz * seqlen, nheads_k, d)
v_ = value.contiguous().reshape(bsz * seqlen, nheads_v, d)
cu = _get_cu_seqlens(q_.device.index, bsz, seqlen)
out = flash_attn_varlen_func_upstream(
q_,
k_,
v_,
cu,
cu,
seqlen,
seqlen,
softmax_scale=self.softmax_scale,
causal=self.causal,
)
return out.reshape(bsz, seqlen, nheads_q, d)
output = flash_attn_func(
q=query, # type: ignore[no-untyped-call]
k=key,

160
python/sglang/multimodal_gen/runtime/layers/triton_ops.py
View File

@@ -129,6 +129,94 @@ def fuse_scale_shift_kernel_blc_opt(
tl.store(y_ptr + x_off, y, mask=mask)
@triton.jit
def fuse_scale_shift_gate_select01_kernel_blc_opt(
x_ptr,
shift0_ptr,
scale0_ptr,
gate0_ptr,
shift1_ptr,
scale1_ptr,
gate1_ptr,
index_ptr,
y_ptr,
gate_out_ptr,
B,
L,
C,
stride_x_b,
stride_x_l,
stride_x_c,
stride_s0_b,
stride_s0_c,
stride_sc0_b,
stride_sc0_c,
stride_g0_b,
stride_g0_c,
stride_s1_b,
stride_s1_c,
stride_sc1_b,
stride_sc1_c,
stride_g1_b,
stride_g1_c,
stride_i_b,
stride_i_l,
stride_go_b,
stride_go_l,
stride_go_c,
BLOCK_L: tl.constexpr,
BLOCK_C: tl.constexpr,
):
pid_l = tl.program_id(0)
pid_c = tl.program_id(1)
pid_b = tl.program_id(2)
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
mask_l = l_offsets < L
mask_c = c_offsets < C
mask = mask_l[:, None] & mask_c[None, :]
x_off = (
pid_b * stride_x_b
+ l_offsets[:, None] * stride_x_l
+ c_offsets[None, :] * stride_x_c
)
x = tl.load(x_ptr + x_off, mask=mask, other=0)
idx_off = pid_b * stride_i_b + l_offsets * stride_i_l
idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None]
s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c
sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c
g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c
s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c
sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c
g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c
shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0)
scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0)
gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0)
shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0)
scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0)
gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0)
shift = tl.where(idx, shift1, shift0)
scale = tl.where(idx, scale1, scale0)
gate = tl.where(idx, gate1, gate0)
y = x * (1 + scale) + shift
tl.store(y_ptr + x_off, y, mask=mask)
go_off = (
pid_b * stride_go_b
+ l_offsets[:, None] * stride_go_l
+ c_offsets[None, :] * stride_go_c
)
tl.store(gate_out_ptr + go_off, gate, mask=mask)
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
@@ -241,6 +329,78 @@ def fuse_scale_shift_kernel(
return output
def fuse_scale_shift_gate_select01_kernel(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
block_l: int = 128,
block_c: int = 128,
):
assert x.is_contiguous()
B, L, C = x.shape
output = torch.empty_like(x)
gate_out = torch.empty_like(x)
if (
scale0.dim() != 2
or shift0.dim() != 2
or gate0.dim() != 2
or scale1.dim() != 2
or shift1.dim() != 2
or gate1.dim() != 2
):
raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
if index.dim() != 2:
raise ValueError("index must be 2D [B, L]")
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
fuse_scale_shift_gate_select01_kernel_blc_opt[grid](
x,
shift0,
scale0,
gate0,
shift1,
scale1,
gate1,
index,
output,
gate_out,
B,
L,
C,
x.stride(0),
x.stride(1),
x.stride(2),
shift0.stride(0),
shift0.stride(1),
scale0.stride(0),
scale0.stride(1),
gate0.stride(0),
gate0.stride(1),
shift1.stride(0),
shift1.stride(1),
scale1.stride(0),
scale1.stride(1),
gate1.stride(0),
gate1.stride(1),
index.stride(0),
index.stride(1),
gate_out.stride(0),
gate_out.stride(1),
gate_out.stride(2),
BLOCK_L=block_l,
BLOCK_C=block_c,
num_warps=4,
num_stages=2,
)
return output, gate_out
@triton.autotune(
configs=[
triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2),

178
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.attention import FeedForward
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -19,7 +20,7 @@ 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_gate_select01_kernel,
fuse_scale_shift_kernel,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
@@ -29,16 +30,24 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
try:
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
except Exception:
apply_rope_with_cos_sin_cache_inplace = None
def _get_qkv_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)
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 attn.added_kv_proj_dim is not None:
txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1)
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
@@ -464,20 +473,27 @@ class QwenImageCrossAttention(nn.Module):
self.parallel_attention = parallel_attention
self.added_kv_proj_dim = added_kv_proj_dim
# Use ReplicatedLinear for fused QKV projections
qkv_dim = num_heads * head_dim * 3
self.to_qkv = ReplicatedLinear(dim, qkv_dim, bias=True)
# Use separate Q/K/V projections
self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
self.inner_kv_dim = self.inner_dim
self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
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:
# Use ReplicatedLinear for added (encoder) QKV projections
self.to_added_qkv = ReplicatedLinear(added_kv_proj_dim, qkv_dim, bias=True)
self.add_q_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=True
)
self.add_k_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=True
)
self.add_v_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)
@@ -545,19 +561,52 @@ class QwenImageCrossAttention(nn.Module):
# 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
)
if apply_rope_with_cos_sin_cache_inplace is None:
raise RuntimeError("flashinfer is required")
if not (
isinstance(image_rotary_emb[0], torch.Tensor)
and image_rotary_emb[0].dim() == 2
):
raise RuntimeError("image_rotary_emb must be cos_sin_cache tensors")
img_cache, txt_cache = image_rotary_emb
def _apply_flashinfer_rope(
q_4d: torch.Tensor, k_4d: torch.Tensor, cache: torch.Tensor
):
bsz, seqlen, nheads, d = q_4d.shape
pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
if bsz == 1:
positions = pos_1d.to(q_4d.device, non_blocking=True)
q2 = q_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
k2 = k_4d.squeeze(0).reshape(seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q2,
key=k2,
head_size=d,
cos_sin_cache=cache,
is_neox=False,
)
return q2.view(1, seqlen, nheads, d), k2.view(1, seqlen, nheads, d)
positions = pos_1d.repeat(bsz).to(q_4d.device, non_blocking=True)
q2 = q_4d.reshape(bsz * seqlen, nheads * d).contiguous()
k2 = k_4d.reshape(bsz * seqlen, nheads * d).contiguous()
apply_rope_with_cos_sin_cache_inplace(
positions=positions,
query=q2,
key=k2,
head_size=d,
cos_sin_cache=cache,
is_neox=False,
)
return q2.view(bsz, seqlen, nheads, d), k2.view(bsz, seqlen, nheads, d)
img_query, img_key = _apply_flashinfer_rope(img_query, img_key, img_cache)
txt_query, txt_key = _apply_flashinfer_rope(txt_query, txt_key, txt_cache)
# Concatenate for joint attention
# Order: [text, image]
@@ -645,32 +694,66 @@ class QwenImageTransformerBlock(nn.Module):
def _modulate(self, x, mod_params, index=None):
shift, scale, gate = mod_params.chunk(3, dim=-1)
if index is not None:
shift_result = shift[index]
scale_result = scale[index]
gate_result = gate[index]
else:
shift_result = shift
scale_result = scale
gate_result = gate[:1].unsqueeze(1)
actual_batch = x.shape[0]
shift0, shift1 = (
shift[:actual_batch],
shift[actual_batch : 2 * actual_batch],
)
scale0, scale1 = (
scale[:actual_batch],
scale[actual_batch : 2 * actual_batch],
)
gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch]
return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result
if x.is_cuda:
if not x.is_contiguous():
x = x.contiguous()
if not index.is_contiguous():
index = index.contiguous()
x, gate_result = fuse_scale_shift_gate_select01_kernel(
x,
scale0=scale0.contiguous(),
shift0=shift0.contiguous(),
gate0=gate0.contiguous(),
scale1=scale1.contiguous(),
shift1=shift1.contiguous(),
gate1=gate1.contiguous(),
index=index,
)
return x, gate_result
else:
mask = (index == 0).unsqueeze(-1)
shift_result = torch.where(
mask, shift0.unsqueeze(1), shift1.unsqueeze(1)
)
scale_result = torch.where(
mask, scale0.unsqueeze(1), scale1.unsqueeze(1)
)
gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1))
return (
fuse_scale_shift_kernel(x, scale_result, shift_result),
gate_result,
)
else:
shift_result = shift.unsqueeze(1)
scale_result = scale.unsqueeze(1)
gate_result = gate.unsqueeze(1)
return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
encoder_hidden_states_mask: torch.Tensor,
temb: torch.Tensor,
temb_img_silu: torch.Tensor,
temb_txt_silu: torch.Tensor,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
modulate_index: Optional[List[int]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# Get modulation parameters for both streams
img_mod_params = self.img_mod(temb) # [B, 6*dim]
if self.zero_cond_t:
temb = torch.chunk(temb, 2, dim=0)[0]
txt_mod_params = self.txt_mod(temb) # [B, 6*dim]
img_mod_params = self.img_mod[1](temb_img_silu) # [B, 6*dim]
txt_mod_params = self.txt_mod[1](temb_txt_silu) # [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]
@@ -880,13 +963,22 @@ class QwenImageTransformer2DModel(CachableDiT):
temb = self.time_text_embed(timestep, hidden_states, additional_t_cond)
temb_img_silu = F.silu(temb)
if self.zero_cond_t:
temb_txt = temb.chunk(2, dim=0)[0]
temb_txt_silu = temb_img_silu.chunk(2, dim=0)[0]
else:
temb_txt = temb
temb_txt_silu = temb_img_silu
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,
temb_img_silu=temb_img_silu,
temb_txt_silu=temb_txt_silu,
image_rotary_emb=image_rotary_emb,
joint_attention_kwargs=attention_kwargs,
modulate_index=modulate_index,
@@ -902,10 +994,8 @@ class QwenImageTransformer2DModel(CachableDiT):
hidden_states
+ controlnet_block_samples[index_block // interval_control]
)
if self.zero_cond_t:
temb = temb.chunk(2, dim=0)[0]
# Use only the image part (hidden_states) from the dual-stream blocks
hidden_states = self.norm_out(hidden_states, temb)
hidden_states = self.norm_out(hidden_states, temb_txt)
output = self.proj_out(hidden_states)
return output