diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index b107243f9..b9e86b6ad 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -72,3 +72,71 @@ class ZImagePipelineConfig(ImagePipelineConfig): def post_denoising_loop(self, latents, batch): bs, channels, num_frames, height, width = latents.shape return latents.view(bs, channels, height, width) + + def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + + axes = [ + torch.arange(x0, x0 + span, dtype=torch.int32, device=device) + for x0, span in zip(start, size) + ] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, dim=-1) + + PATCH_SIZE = 2 + F_PATCH_SIZE = 1 + SEQ_MULTI_OF = 32 + + cap_ori_len = prompt_embeds.size(0) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + cap_padded_pos_ids = create_coordinate_grid( + size=(cap_ori_len + cap_padding_len, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2) + + C = self.dit_config.num_channels_latents + F = 1 + H = height // self.vae_config.arch_config.spatial_compression_ratio + W = width // self.vae_config.arch_config.spatial_compression_ratio + + pH, pW = PATCH_SIZE, PATCH_SIZE + pF = F_PATCH_SIZE + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + image_ori_len = F_tokens * H_tokens * W_tokens + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + + image_ori_pos_ids = create_coordinate_grid( + size=(F_tokens, H_tokens, W_tokens), + start=(cap_ori_len + cap_padding_len + 1, 0, 0), + device=device, + ).flatten(0, 2) + image_padding_pos_ids = ( + create_coordinate_grid( + size=(1, 1, 1), + start=(0, 0, 0), + device=device, + ) + .flatten(0, 2) + .repeat(image_padding_len, 1) + ) + image_padded_pos_ids = torch.cat( + [image_ori_pos_ids, image_padding_pos_ids], dim=0 + ) + cap_freqs_cis = rotary_emb(cap_padded_pos_ids) + x_freqs_cis = rotary_emb(image_padded_pos_ids) + return (cap_freqs_cis, x_freqs_cis) + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "freqs_cis": self.get_freqs_cis( + batch.prompt_embeds[0], + batch.width, + batch.height, + device, + rotary_emb, + batch, + ), + } diff --git a/python/sglang/multimodal_gen/configs/sample/zimage.py b/python/sglang/multimodal_gen/configs/sample/zimage.py index 1ae8bc0de..f4530e977 100644 --- a/python/sglang/multimodal_gen/configs/sample/zimage.py +++ b/python/sglang/multimodal_gen/configs/sample/zimage.py @@ -10,6 +10,13 @@ from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @dataclass class ZImageSamplingParams(SamplingParams): num_inference_steps: int = 9 + + num_frames: int = 1 + negative_prompt: str = None + # height: int = 720 + # width: int = 1280 + # fps: int = 24 + guidance_scale: float = 0.0 teacache_params: TeaCacheParams = field( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index ab9af9114..bdc38d52d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -4,12 +4,12 @@ from typing import Any, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F -from torch.nn.utils.rnn import pad_sequence from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -150,13 +150,12 @@ class ZImageAttention(nn.Module): # Apply RoPE def apply_rotary_emb( - x_in: torch.Tensor, freqs_cis: torch.Tensor - ) -> torch.Tensor: - with torch.amp.autocast("cuda", enabled=False): - x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.unsqueeze(2) - x_out = torch.view_as_real(x * freqs_cis).flatten(3) - return x_out.type_as(x_in) # todo + x_in: torch.Tensor, + freqs_cis: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + cos, sin = freqs_cis + x_out = _apply_rotary_emb(x_in, cos, sin, is_neox_style=False) + return x_out if freqs_cis is not None: q = apply_rotary_emb(q, freqs_cis) @@ -211,7 +210,6 @@ class ZImageTransformerBlock(nn.Module): def forward( self, x: torch.Tensor, - attn_mask: torch.Tensor, freqs_cis: Tuple[torch.Tensor, torch.Tensor], adaln_input: Optional[torch.Tensor] = None, ): @@ -288,12 +286,15 @@ class RopeEmbedder: assert len(axes_dims) == len( axes_lens ), "axes_dims and axes_lens must have the same length" - self.freqs_cis = None + + self.cos_cached = None + self.sin_cached = None @staticmethod - def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + def precompute_freqs(dim: List[int], end: List[int], theta: float = 256.0): with torch.device("cpu"): - freqs_cis = [] + cos_list = [] + sin_list = [] for i, (d, e) in enumerate(zip(dim, end)): freqs = 1.0 / ( theta @@ -301,33 +302,43 @@ class RopeEmbedder: ) timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) freqs = torch.outer(timestep, freqs).float() - freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to( - torch.complex64 - ) # complex64 - freqs_cis.append(freqs_cis_i) - return freqs_cis + cos_list.append(torch.cos(freqs)) + sin_list.append(torch.sin(freqs)) - def __call__(self, ids: torch.Tensor): + return cos_list, sin_list + + def __call__(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + ids: [batch, len(axes_dims)] or [seq_len, len(axes_dims)] + Returns: + cos: [batch/seq, head_dim // 2] + sin: [batch/seq, head_dim // 2] + """ assert ids.ndim == 2 assert ids.shape[-1] == len(self.axes_dims) device = ids.device - if self.freqs_cis is None: - self.freqs_cis = self.precompute_freqs_cis( + if self.cos_cached is None: + self.cos_cached, self.sin_cached = self.precompute_freqs( self.axes_dims, self.axes_lens, theta=self.theta ) - self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + self.cos_cached = [c.to(device) for c in self.cos_cached] + self.sin_cached = [s.to(device) for s in self.sin_cached] else: - # Ensure freqs_cis are on the same device as ids - if self.freqs_cis[0].device != device: - self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + if self.cos_cached[0].device != device: + self.cos_cached = [c.to(device) for c in self.cos_cached] + self.sin_cached = [s.to(device) for s in self.sin_cached] - result = [] + cos_out = [] + sin_out = [] for i in range(len(self.axes_dims)): index = ids[:, i] - result.append(self.freqs_cis[i][index]) - return torch.cat(result, dim=-1) + cos_out.append(self.cos_cached[i][index]) + sin_out.append(self.sin_cached[i][index]) + + return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1) class ZImageTransformer2DModel(CachableDiT): @@ -435,7 +446,7 @@ class ZImageTransformer2DModel(CachableDiT): self.axes_dims = arch_config.axes_dims self.axes_lens = arch_config.axes_lens - self.rope_embedder = RopeEmbedder( + self.rotary_emb = RopeEmbedder( theta=self.rope_theta, axes_dims=self.axes_dims, axes_lens=self.axes_lens ) @@ -477,114 +488,64 @@ class ZImageTransformer2DModel(CachableDiT): patch_size: int, f_patch_size: int, ): + assert len(all_image) == len(all_cap_feats) == 1 + + image = all_image[0] # C, F, H, W + cap_feat = all_cap_feats[0] # L, D pH = pW = patch_size pF = f_patch_size - device = all_image[0].device + device = image.device all_image_out = [] all_image_size = [] - all_image_pos_ids = [] - all_image_pad_mask = [] - all_cap_pos_ids = [] - all_cap_pad_mask = [] all_cap_feats_out = [] - for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): - ### Process Caption - cap_ori_len = len(cap_feat) - cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF - # padded position ids - cap_padded_pos_ids = self.create_coordinate_grid( - size=(cap_ori_len + cap_padding_len, 1, 1), - start=(1, 0, 0), - device=device, - ).flatten(0, 2) - all_cap_pos_ids.append(cap_padded_pos_ids) - # pad mask - all_cap_pad_mask.append( - torch.cat( - [ - torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), - torch.ones((cap_padding_len,), dtype=torch.bool, device=device), - ], - dim=0, - ) - ) - # padded feature - cap_padded_feat = torch.cat( - [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], - dim=0, - ) - all_cap_feats_out.append(cap_padded_feat) + # ------------ Process Caption ------------ + cap_ori_len = cap_feat.size(0) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF - ### Process Image - C, F, H, W = image.size() - all_image_size.append((F, H, W)) - F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + # padded feature + cap_padded_feat = torch.cat( + [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], + dim=0, + ) + all_cap_feats_out.append(cap_padded_feat) - image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) - # "c f pf h ph w pw -> (f h w) (pf ph pw c)" - image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( - F_tokens * H_tokens * W_tokens, pF * pH * pW * C - ) + # ------------ Process Image ------------ + C, F, H, W = image.size() + all_image_size.append((F, H, W)) - image_ori_len = len(image) - image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + # "c f pf h ph w pw -> (f h w) (pf ph pw c)" + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( + F_tokens * H_tokens * W_tokens, pF * pH * pW * C + ) + image_ori_len = image.size(0) + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF - image_ori_pos_ids = self.create_coordinate_grid( - size=(F_tokens, H_tokens, W_tokens), - start=(cap_ori_len + cap_padding_len + 1, 0, 0), - device=device, - ).flatten(0, 2) - image_padding_pos_ids = ( - self.create_coordinate_grid( - size=(1, 1, 1), - start=(0, 0, 0), - device=device, - ) - .flatten(0, 2) - .repeat(image_padding_len, 1) - ) - image_padded_pos_ids = torch.cat( - [image_ori_pos_ids, image_padding_pos_ids], dim=0 - ) - all_image_pos_ids.append(image_padded_pos_ids) - # pad mask - all_image_pad_mask.append( - torch.cat( - [ - torch.zeros((image_ori_len,), dtype=torch.bool, device=device), - torch.ones( - (image_padding_len,), dtype=torch.bool, device=device - ), - ], - dim=0, - ) - ) - # padded feature - image_padded_feat = torch.cat( - [image, image[-1:].repeat(image_padding_len, 1)], dim=0 - ) - all_image_out.append(image_padded_feat) + # padded feature + image_padded_feat = torch.cat( + [image, image[-1:].repeat(image_padding_len, 1)], + dim=0, + ) + all_image_out.append(image_padded_feat) return ( all_image_out, all_cap_feats_out, all_image_size, - all_image_pos_ids, - all_cap_pos_ids, - all_image_pad_mask, - all_cap_pad_mask, ) def forward( self, hidden_states: List[torch.Tensor], - timestep, encoder_hidden_states: List[torch.Tensor], + timestep, guidance=0, patch_size=2, f_patch_size=1, + freqs_cis=None, **kwargs, ): assert patch_size in self.all_patch_size @@ -594,129 +555,43 @@ class ZImageTransformer2DModel(CachableDiT): cap_feats = encoder_hidden_states timestep = 1000.0 - timestep t = timestep - bsz = len(x) + bsz = 1 device = x[0].device - # t = t * self.t_scale t = self.t_embedder(t) - adaln_input = t.type_as(x) ( x, cap_feats, x_size, - x_pos_ids, - cap_pos_ids, - x_inner_pad_mask, - cap_inner_pad_mask, ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) - # x embed & refine - x_item_seqlens = [len(_) for _ in x] - assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) x = torch.cat(x, dim=0) x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) - x[torch.cat(x_inner_pad_mask)] = self.x_pad_token - x = list(x.split(x_item_seqlens, dim=0)) - x_freqs_cis = list( - self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0) - ) + x_freqs_cis = freqs_cis[1] - # RoPE returns (cos, sin) now - # x_pos_ids_cat = torch.cat(x_pos_ids, dim=0) - # x_cos, x_sin = self.rope_embedder(x_pos_ids_cat) - - # x_cos_list = list(x_cos.split(x_item_seqlens, dim=0)) - # x_sin_list = list(x_sin.split(x_item_seqlens, dim=0)) - - x = pad_sequence(x, batch_first=True, padding_value=0.0) - x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - # x_cos = pad_sequence(x_cos_list, batch_first=True, padding_value=0.0) - # x_sin = pad_sequence(x_sin_list, batch_first=True, padding_value=0.0) - # B, T, D_half = x_cos.shape # D_half = 64 - - # x_cos_triton = x_cos.reshape(B * T, D_half).contiguous() # [B*T, 64] - # x_sin_triton = x_sin.reshape(B * T, D_half).contiguous() # [B*T, 64] - - x_attn_mask = torch.zeros( - (bsz, x_max_item_seqlen), dtype=torch.bool, device=device - ) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Refiner logic + x = x.unsqueeze(0) + x_freqs_cis = x_freqs_cis for layer in self.noise_refiner: - x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) - - # cap embed & refine - cap_item_seqlens = [len(_) for _ in cap_feats] - assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) - cap_max_item_seqlen = max(cap_item_seqlens) + x = layer(x, x_freqs_cis, adaln_input) cap_feats = torch.cat(cap_feats, dim=0) - # cap_embedder is Sequential with ReplicatedLinear. - # We need to handle this if ReplicatedLinear returns tuple. - # In __init__, cap_embedder = Sequential(RMSNorm, ReplicatedLinear). - # RMSNorm returns Tensor. ReplicatedLinear returns (Tensor, Gathered). - # Sequential returns (Tensor, Gathered). - # So we need to unpack. - cap_feats_out = self.cap_embedder(cap_feats) - if isinstance(cap_feats_out, tuple): - cap_feats = cap_feats_out[0] - else: - cap_feats = cap_feats_out + cap_feats, _ = self.cap_embedder(cap_feats) - cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token - cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) - - cap_freqs_cis = list( - self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split( - cap_item_seqlens, dim=0 - ) - ) - - cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) - cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) - - cap_attn_mask = torch.zeros( - (bsz, cap_max_item_seqlen), dtype=torch.bool, device=device - ) - for i, seq_len in enumerate(cap_item_seqlens): - cap_attn_mask[i, :seq_len] = 1 + cap_freqs_cis = freqs_cis[0] + cap_feats = cap_feats.unsqueeze(0) for layer in self.context_refiner: - cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) + cap_feats = layer(cap_feats, cap_freqs_cis) - # unified - unified = [] - unified_freqs_cis = [] - - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) - unified_freqs_cis.append( - torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]]) - ) - - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] - assert unified_item_seqlens == [len(_) for _ in unified] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_freqs_cis = pad_sequence( - unified_freqs_cis, batch_first=True, padding_value=0.0 + unified = torch.cat([x, cap_feats], dim=1) + unified_freqs_cis = ( + torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=0), + torch.cat([x_freqs_cis[1], cap_freqs_cis[1]], dim=0), ) - unified_attn_mask = torch.zeros( - (bsz, unified_max_item_seqlen), dtype=torch.bool, device=device - ) - for i, seq_len in enumerate(unified_item_seqlens): - unified_attn_mask[i, :seq_len] = 1 - for layer in self.layers: - unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) + unified = layer(unified, unified_freqs_cis, adaln_input) unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( unified, adaln_input diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 8febe38dc..e74c755f1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -850,18 +850,18 @@ class DenoisingStage(PipelineStage): # Predict noise residual attn_metadata = self._build_attn_metadata(i, batch, server_args) noise_pred = self._predict_noise_with_cfg( - current_model, - latent_model_input, - timestep, - batch, - i, - attn_metadata, - target_dtype, - current_guidance_scale, - image_kwargs, - pos_cond_kwargs, - neg_cond_kwargs, - server_args, + current_model=current_model, + latent_model_input=latent_model_input, + timestep=timestep, + batch=batch, + timestep_index=i, + attn_metadata=attn_metadata, + target_dtype=target_dtype, + current_guidance_scale=current_guidance_scale, + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + neg_cond_kwargs=neg_cond_kwargs, + server_args=server_args, guidance=guidance, latents=latents, )