From a1e9b4edfa34640d7757102d6e764e9c1570f597 Mon Sep 17 00:00:00 2001 From: chhnb <44352333+chhnb@users.noreply.github.com> Date: Fri, 26 Dec 2025 14:57:01 +0800 Subject: [PATCH] diffusion: support Qwen-Image-Layered (#15817) Co-authored-by: yhyang201 --- .../configs/models/vaes/qwenimage.py | 2 +- .../configs/pipeline_configs/base.py | 3 + .../configs/pipeline_configs/qwen_image.py | 88 +++ .../configs/sample/qwenimage.py | 14 + .../configs/sample/sampling_params.py | 8 +- python/sglang/multimodal_gen/registry.py | 8 + .../entrypoints/diffusion_generator.py | 7 +- .../runtime/entrypoints/utils.py | 11 +- .../runtime/models/dits/qwen_image.py | 240 +++++++- .../models/model_stages/qwen_image_layered.py | 528 ++++++++++++++++++ .../models/vaes/autoencoder_kl_qwenimage.py | 14 +- .../runtime/pipelines/qwen_image.py | 63 ++- 12 files changed, 962 insertions(+), 24 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/models/model_stages/qwen_image_layered.py diff --git a/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py b/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py index d08ab5da9..af9fa9d2a 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py @@ -20,7 +20,7 @@ class QwenImageVAEArchConfig(VAEArchConfig): dropout: float = 0.0 is_residual: bool = False - in_channels: int = 3 + input_channels: int = 3 out_channels: int = 3 patch_size: int | None = None scale_factor_temporal: int = 4 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 4a2d22ead..570da92c8 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -299,6 +299,9 @@ class PipelineConfig: return shape + def allow_set_num_frames(self): + return False + def get_decode_scale_and_shift(self, device, dtype, vae): vae_arch_config = self.vae_config.arch_config scaling_factor = getattr(vae_arch_config, "scaling_factor", None) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 54066eaf5..042013495 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -12,6 +12,7 @@ from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConf from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, + maybe_unpad_latents, shard_rotary_emb_for_sp, ) from sglang.multimodal_gen.runtime.models.vision_utils import resize @@ -466,3 +467,90 @@ class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig): "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)), "img_shapes": img_shapes, } + + +@dataclass +class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig): + resolution: int = 640 # TODO: allow user to set resolution + vae_precision: str = "bf16" + + def _prepare_edit_cond_kwargs( + self, batch, prompt_embeds, rotary_emb, device, dtype + ): + batch_size = batch.latents.shape[0] + assert batch_size == 1 + height = batch.height + width = batch.width + image_size = batch.original_condition_image_size + + vae_scale_factor = self.get_vae_scale_factor() + + img_shapes = batch.img_shapes + txt_seq_lens = batch.txt_seq_lens + + (img_cos, img_sin), (txt_cos, txt_sin) = ( + 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) + ) + 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 + ) + + return { + "txt_seq_lens": txt_seq_lens, + "img_shapes": img_shapes, + "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)), + "additional_t_cond": torch.tensor([0], device=device, dtype=torch.long), + } + + def _unpad_and_unpack_latents(self, latents, batch): + vae_scale_factor = self.vae_config.arch_config.vae_scale_factor + channels = self.dit_config.arch_config.in_channels + batch_size = latents.shape[0] + layers = batch.num_frames + + height = 2 * (int(batch.height) // (vae_scale_factor * 2)) + width = 2 * (int(batch.width) // (vae_scale_factor * 2)) + + latents = maybe_unpad_latents(latents, batch) + latents = latents.view( + batch_size, layers + 1, height // 2, width // 2, channels // 4, 2, 2 + ) + latents = latents.permute(0, 1, 4, 2, 5, 3, 6) + + latents = latents.reshape( + batch_size, layers + 1, channels // (2 * 2), height, width + ) + latents = latents.permute(0, 2, 1, 3, 4) # (b, c, f, h, w) + return latents, batch_size, channels, height, width + + def allow_set_num_frames(self): + return True + + def post_denoising_loop(self, latents, batch): + # unpack latents for qwen-image + ( + latents, + batch_size, + channels, + height, + width, + ) = self._unpad_and_unpack_latents(latents, batch) + b, c, f, h, w = latents.shape + latents = latents[:, :, 1:] # remove the first frame as it is the origin input + latents = latents.permute(0, 2, 1, 3, 4).view(-1, c, 1, h, w) + # latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width) + return latents diff --git a/python/sglang/multimodal_gen/configs/sample/qwenimage.py b/python/sglang/multimodal_gen/configs/sample/qwenimage.py index c043f328b..c24c1dd25 100644 --- a/python/sglang/multimodal_gen/configs/sample/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/sample/qwenimage.py @@ -21,3 +21,17 @@ class QwenImageEditPlusSamplingParams(QwenImageSamplingParams): guidance_scale: float = 4.0 # true_cfg_scale: float = 4.0 num_inference_steps: int = 40 + + +@dataclass +class QwenImageLayeredSamplingParams(QwenImageSamplingParams): + # num_frames: int = 4 + height: int = 640 + width: int = 640 + prompt: str = " " + negative_prompt: str = " " + + guidance_scale: float = 4.0 + num_inference_steps: int = 50 + cfg_normalize: bool = True + use_en_prompt: bool = True diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 81a442a90..2adc7c94f 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -69,7 +69,7 @@ class DataType(Enum): def get_default_extension(self) -> str: if self == DataType.IMAGE: - return "jpg" + return "png" else: return "mp4" @@ -279,8 +279,10 @@ class SamplingParams: if pipeline_config.task_type.is_image_gen(): # settle num_frames - logger.debug(f"num_frames set to 1 for image generation model") - self.num_frames = 1 + if not server_args.pipeline_config.allow_set_num_frames(): + logger.debug(f"num_frames set to 1 for image generation model") + self.num_frames = 1 + elif self.adjust_frames: # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below. # If we apply it after, adjust_num_frames might modify the frame count diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 3797fa70d..7d256c0b9 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -30,6 +30,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineCon from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, QwenImageEditPlusPipelineConfig, + QwenImageLayeredPipelineConfig, QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( @@ -46,6 +47,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( ) from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageEditPlusSamplingParams, + QwenImageLayeredSamplingParams, QwenImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams @@ -443,5 +445,11 @@ def _register_configs(): hf_model_paths=["Qwen/Qwen-Image-Edit-2511"], ) + register_configs( + sampling_param_cls=QwenImageLayeredSamplingParams, + pipeline_config_cls=QwenImageLayeredPipelineConfig, + hf_model_paths=["Qwen/Qwen-Image-Layered"], + ) + _register_configs() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 65ac0abcb..1b1cb8d16 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -183,14 +183,13 @@ class DiffGenerator: raise ValueError(f"No prompts found in file: {prompt_txt_path}") logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path) - elif prompt is not None: + else: + if prompt is None: + prompt = " " if isinstance(prompt, str): prompts.append(prompt) elif isinstance(prompt, list): prompts.extend(prompt) - else: - raise ValueError("Either prompt or prompt_txt must be provided") - sampling_params = SamplingParams.from_user_sampling_params_args( self.server_args.model_path, server_args=self.server_args, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 2d8af480e..b2166ebef 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -93,7 +93,16 @@ def post_process_sample( ) else: quality = 75 - imageio.imwrite(save_file_path, frames[0], quality=quality) + if len(frames) > 1: + for i, image in enumerate(frames): + parts = save_file_path.rsplit(".", 1) + if len(parts) == 2: + indexed_path = f"{parts[0]}_{i}.{parts[1]}" + else: + indexed_path = f"{save_file_path}_{i}" + imageio.imwrite(indexed_path, image, quality=quality) + else: + imageio.imwrite(save_file_path, frames[0], quality=quality) logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}") else: logger.info(f"No output path provided, output not saved") diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 9a17b6bd3..169055a7d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -44,7 +44,7 @@ def _get_qkv_projections( class QwenTimestepProjEmbeddings(nn.Module): - def __init__(self, embedding_dim): + def __init__(self, embedding_dim, use_additional_t_cond=False): super().__init__() self.time_proj = Timesteps( @@ -53,14 +53,25 @@ class QwenTimestepProjEmbeddings(nn.Module): self.timestep_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) + self.use_additional_t_cond = use_additional_t_cond + if use_additional_t_cond: + self.addition_t_embedding = nn.Embedding(2, embedding_dim) - def forward(self, timestep, hidden_states): + def forward(self, timestep, hidden_states, addition_t_cond=None): timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder( timesteps_proj.to(dtype=hidden_states.dtype) ) # (N, D) conditioning = timesteps_emb + if self.use_additional_t_cond: + if addition_t_cond is None: + raise ValueError( + "When additional_t_cond is True, addition_t_cond must be provided." + ) + addition_t_emb = self.addition_t_embedding(addition_t_cond) + addition_t_emb = addition_t_emb.to(dtype=hidden_states.dtype) + conditioning = conditioning + addition_t_emb return conditioning @@ -232,6 +243,200 @@ class QwenEmbedRope(nn.Module): return freqs.clone().contiguous() +class QwenEmbedLayer3DRope(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.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, txt_seq_lens, device): + """ + Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args: + txt_length: [bs] a list of 1 integers representing the length of the text + """ + + # 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 + layer_num = len(video_fhw) - 1 + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + if idx != layer_num: + video_freq = self._compute_video_freqs(frame, height, width, idx) + else: + ### For the condition image, we set the layer index to -1 + video_freq = self._compute_condition_freqs(frame, height, width) + 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_vid_index = max(max_vid_index, layer_num) + 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) + + return vid_freqs, txt_freqs + + @functools.lru_cache(maxsize=None) + def _compute_video_freqs(self, frame, height, width, idx=0): + 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() + + @functools.lru_cache(maxsize=None) + def _compute_condition_freqs(self, frame, height, width): + 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_neg[0][-1:].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): def __init__( self, @@ -553,17 +758,31 @@ class QwenImageTransformer2DModel(CachableDiT): 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 - zero_cond_t = getattr(config.arch_config, "zero_cond_t", False) + self.zero_cond_t = getattr(config.arch_config, "zero_cond_t", False) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim - self.zero_cond_t = zero_cond_t - self.rotary_emb = QwenEmbedRope( - theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True + self.use_additional_t_cond: bool = getattr( + config.arch_config, "use_additional_t_cond", False + ) # For qwen-image-layered now + self.use_layer3d_rope: bool = getattr( + config.arch_config, "use_layer3d_rope", False + ) # For qwen-image-layered now + + if not self.use_layer3d_rope: + self.rotary_emb = QwenEmbedRope( + theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True + ) + else: + self.rotary_emb = QwenEmbedLayer3DRope( + theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True + ) + + self.time_text_embed = QwenTimestepProjEmbeddings( + embedding_dim=self.inner_dim, + use_additional_t_cond=self.use_additional_t_cond, ) - 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) @@ -575,7 +794,7 @@ class QwenImageTransformer2DModel(CachableDiT): dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, - zero_cond_t=zero_cond_t, + zero_cond_t=self.zero_cond_t, ) for _ in range(num_layers) ] @@ -597,6 +816,7 @@ class QwenImageTransformer2DModel(CachableDiT): img_shapes: Optional[List[Tuple[int, int, int]]] = None, txt_seq_lens: Optional[List[int]] = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] = None, + additional_t_cond: Optional[torch.Tensor] = None, guidance: torch.Tensor = None, # TODO: this should probably be removed attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_block_samples=None, @@ -658,7 +878,7 @@ class QwenImageTransformer2DModel(CachableDiT): 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) + temb = self.time_text_embed(timestep, hidden_states, additional_t_cond) image_rotary_emb = freqs_cis for index_block, block in enumerate(self.transformer_blocks): diff --git a/python/sglang/multimodal_gen/runtime/models/model_stages/qwen_image_layered.py b/python/sglang/multimodal_gen/runtime/models/model_stages/qwen_image_layered.py new file mode 100644 index 000000000..999aebd58 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/model_stages/qwen_image_layered.py @@ -0,0 +1,528 @@ +import inspect +import math +from typing import List, Optional, Union + +import numpy as np +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.vision_utils import load_image +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus.calculate_dimensions +def calculate_dimensions(target_area, ratio): + width = math.sqrt(target_area * ratio) + height = width / ratio + + width = round(width / 32) * 32 + height = round(height / 32) * 32 + + return width, height + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, + generator: Optional[torch.Generator] = None, + sample_mode: str = "sample", +): + if sample_mode == "sample": + return encoder_output.sample(generator) + elif sample_mode == "argmax": + return encoder_output.mode() + else: + return encoder_output + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError( + "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" + ) + if timesteps is not None: + accepts_timesteps = "timesteps" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys() + ) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class QwenImageLayeredBeforeDenoisingStage(PipelineStage): + def __init__( + self, vae, tokenizer, processor, transformer, scheduler, model_path + ) -> None: + self.vae = vae.to(torch.bfloat16) + from transformers import Qwen2_5_VLForConditionalGeneration + + self.text_encoder = ( + Qwen2_5_VLForConditionalGeneration.from_pretrained( + model_path, subfolder="text_encoder" + ) + .to(get_local_torch_device()) + .to(torch.bfloat16) + ) + self.tokenizer = tokenizer + self.processor = processor + self.transformer = transformer + self.scheduler = scheduler + + self.vae_scale_factor = ( + 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + ) + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor * 2 + ) + self.vl_processor = processor + self.tokenizer_max_length = 1024 + self.latent_channels = self.vae.z_dim if getattr(self, "vae", None) else 16 + + self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + self.prompt_template_encode_start_idx = 34 + self.image_caption_prompt_cn = """<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n# 图像标注器\n你是一个专业的图像标注器。请基于输入图像，撰写图注:\n1. +使用自然、描述性的语言撰写图注，不要使用结构化形式或富文本形式。\n2. 通过加入以下内容，丰富图注细节：\n - 对象的属性：如数量、颜色、形状、大小、位置、材质、状态、动作等\n - +对象间的视觉关系：如空间关系、功能关系、动作关系、从属关系、比较关系、因果关系等\n - 环境细节：例如天气、光照、颜色、纹理、气氛等\n - 文字内容：识别图像中清晰可见的文字，不做翻译和解释，用引号在图注中强调\n3. +保持真实性与准确性：\n - 不要使用笼统的描述\n - +描述图像中所有可见的信息，但不要加入没有在图像中出现的内容\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n""" + self.image_caption_prompt_en = """<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n# Image Annotator\nYou are a professional +image annotator. Please write an image caption based on the input image:\n1. Write the caption using natural, +descriptive language without structured formats or rich text.\n2. Enrich caption details by including: \n - Object +attributes, such as quantity, color, shape, size, material, state, position, actions, and so on\n - Vision Relations +between objects, such as spatial relations, functional relations, possessive relations, attachment relations, action +relations, comparative relations, causal relations, and so on\n - Environmental details, such as weather, lighting, +colors, textures, atmosphere, and so on\n - Identify the text clearly visible in the image, without translation or +explanation, and highlight it in the caption with quotation marks\n3. Maintain authenticity and accuracy:\n - Avoid +generalizations\n - Describe all visible information in the image, while do not add information not explicitly shown in +the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n""" + self.default_sample_size = 128 + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + + return split_result + + def get_image_caption(self, prompt_image, use_en_prompt=True, device=None): + if use_en_prompt: + prompt = self.image_caption_prompt_en + else: + prompt = self.image_caption_prompt_cn + model_inputs = self.vl_processor( + text=prompt, + images=prompt_image, + padding=True, + return_tensors="pt", + ).to(device) + with set_forward_context(current_timestep=0, attn_metadata=None): + generated_ids = self.text_encoder.generate( + **model_inputs, max_new_tokens=512 + ) + generated_ids_trimmed = [ + out_ids[len(in_ids) :] + for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids) + ] + output_text = self.vl_processor.batch_decode( + generated_ids_trimmed, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + )[0] + return output_text.strip() + + def _get_qwen_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + template = self.prompt_template_encode + drop_idx = self.prompt_template_encode_start_idx + txt = [template.format(e) for e in prompt] + txt_tokens = self.tokenizer( + txt, + padding=True, + return_tensors="pt", + ).to(device) + encoder_hidden_states = self.text_encoder( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask, + output_hidden_states=True, + ) + hidden_states = encoder_hidden_states.hidden_states[-1] + split_hidden_states = self._extract_masked_hidden( + hidden_states, txt_tokens.attention_mask + ) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [ + torch.ones(e.size(0), dtype=torch.long, device=e.device) + for e in split_hidden_states + ] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack( + [ + torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) + for u in split_hidden_states + ] + ) + encoder_attention_mask = torch.stack( + [ + torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) + for u in attn_mask_list + ] + ) + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + return prompt_embeds, encoder_attention_mask + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width, layers): + latents = latents.view( + batch_size, layers, num_channels_latents, height // 2, 2, width // 2, 2 + ) + latents = latents.permute(0, 1, 3, 5, 2, 4, 6) + latents = latents.reshape( + batch_size, layers * (height // 2) * (width // 2), num_channels_latents * 4 + ) + + return latents + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + max_sequence_length: int = 1024, + ): + r""" + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, device + ) + + prompt_embeds = prompt_embeds[:, :max_sequence_length] + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(num_images_per_prompt, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + self.vae = self.vae.to(get_local_torch_device()) + if isinstance(generator, list): + image_latents = [ + retrieve_latents( + self.vae.encode(image[i : i + 1]), + generator=generator[i], + sample_mode="argmax", + ) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents( + self.vae.encode(image), generator=generator, sample_mode="argmax" + ) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + self.vae.to("cpu") + return image_latents + + def prepare_latents( + self, + image, + batch_size, + num_channels_latents, + height, + width, + layers, + dtype, + device, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + shape = ( + batch_size, + layers + 1, + num_channels_latents, + height, + width, + ) ### the generated first image is combined image + + image_latents = None + if image is not None: + image = image.to(device=device, dtype=dtype) + if image.shape[1] != self.latent_channels: + image_latents = self._encode_vae_image(image=image, generator=generator) + else: + image_latents = image + if ( + batch_size > image_latents.shape[0] + and batch_size % image_latents.shape[0] == 0 + ): + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat( + [image_latents] * additional_image_per_prompt, dim=0 + ) + elif ( + batch_size > image_latents.shape[0] + and batch_size % image_latents.shape[0] != 0 + ): + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + image_latent_height, image_latent_width = image_latents.shape[3:] + image_latents = image_latents.permute( + 0, 2, 1, 3, 4 + ) # (b, c, f, h, w) -> (b, f, c, h, w) + image_latents = self._pack_latents( + image_latents, + batch_size, + num_channels_latents, + image_latent_height, + image_latent_width, + 1, + ) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + latents = randn_tensor( + shape, generator=generator, device=device, dtype=dtype + ) + latents = self._pack_latents( + latents, batch_size, num_channels_latents, height, width, layers + 1 + ) + else: + latents = latents.to(device=device, dtype=dtype) + + return latents, image_latents + + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + use_en_prompt = True + device = get_local_torch_device() + layers = batch.num_frames + num_inference_steps = batch.num_inference_steps + generator = batch.generator + + assert batch.image_path is not None + image = load_image(batch.image_path[0]) + image = image.convert("RGBA") + image_size = image.size + resolution = 640 # TODO: support user-specified resolution + calculated_width, calculated_height = calculate_dimensions( + resolution * resolution, image_size[0] / image_size[1] + ) + + height = calculated_height + width = calculated_width + + multiple_of = self.vae_scale_factor * 2 + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + + # if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): + image = self.image_processor.resize(image, calculated_height, calculated_width) + prompt_image = image + image = self.image_processor.preprocess( + image, calculated_height, calculated_width + ) + image = image.unsqueeze(2) + image = image.to(dtype=torch.bfloat16) + + prompt = self.get_image_caption( + prompt_image, use_en_prompt=use_en_prompt, device=device + ) + + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + device=device, + ) + + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + prompt=batch.negative_prompt, + device=device, + ) + + num_channels_latents = self.transformer.config.in_channels // 4 + latents, image_latents = self.prepare_latents( + image, + 1, + num_channels_latents, + height, + width, + layers, + prompt_embeds.dtype, + device, + generator, + ) + img_shapes = [ + [ + *[ + ( + 1, + height // self.vae_scale_factor // 2, + width // self.vae_scale_factor // 2, + ) + for _ in range(layers + 1) + ], + ( + 1, + calculated_height // self.vae_scale_factor // 2, + calculated_width // self.vae_scale_factor // 2, + ), + ] + ] + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 0, num_inference_steps + 1)[:-1] + image_seq_len = latents.shape[1] + base_seqlen = 256 * 256 / 16 / 16 + mu = (image_latents.shape[1] / base_seqlen) ** 0.5 + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + + txt_seq_lens = ( + prompt_embeds_mask.sum(dim=1).tolist() + if prompt_embeds_mask is not None + else None + ) + negative_txt_seq_lens = ( + negative_prompt_embeds_mask.sum(dim=1).tolist() + if negative_prompt_embeds_mask is not None + else None + ) + is_rgb = torch.tensor([0]).to(device=device, dtype=torch.long) + + batch.prompt_embeds = [prompt_embeds] + batch.prompt_embeds_mask = [prompt_embeds_mask] + batch.negative_prompt_embeds = [negative_prompt_embeds] + batch.negative_prompt_embeds_mask = [negative_prompt_embeds_mask] + batch.latents = latents + batch.image_latent = image_latents + batch.num_inference_steps = num_inference_steps + batch.sigmas = sigmas.tolist() # Convert numpy array to list for validation + batch.generator = torch.manual_seed(0) + batch.original_condition_image_size = image_size + batch.raw_latent_shape = latents.shape + batch.txt_seq_lens = txt_seq_lens + batch.img_shapes = img_shapes + + return batch diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py index e860b313b..01d3c0114 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py @@ -445,6 +445,7 @@ class QwenImageEncoder3d(nn.Module): temperal_downsample=[True, True, False], dropout=0.0, non_linearity: str = "silu", + input_channels: int = 3, ): super().__init__() # dim = config.arch_config.dim @@ -468,7 +469,7 @@ class QwenImageEncoder3d(nn.Module): scale = 1.0 # init block - self.conv_in = QwenImageCausalConv3d(3, dims[0], 3, padding=1) + self.conv_in = QwenImageCausalConv3d(input_channels, dims[0], 3, padding=1) # downsample blocks self.down_blocks = nn.ModuleList([]) @@ -649,6 +650,7 @@ class QwenImageDecoder3d(nn.Module): temperal_upsample=[False, True, True], dropout=0.0, non_linearity: str = "silu", + input_channels=3, ): super().__init__() self.dim = dim @@ -701,7 +703,7 @@ class QwenImageDecoder3d(nn.Module): # output blocks self.norm_out = QwenImageRMS_norm(out_dim, images=False) - self.conv_out = QwenImageCausalConv3d(out_dim, 3, 3, padding=1) + self.conv_out = QwenImageCausalConv3d(out_dim, input_channels, 3, padding=1) self.gradient_checkpointing = False @@ -783,15 +785,19 @@ class AutoencoderKLQwenImage(nn.Module): self.z_dim = z_dim self.temperal_downsample = temperal_downsample self.temperal_upsample = temperal_downsample[::-1] + self.input_channels = config.arch_config.input_channels + self.latents_mean = config.arch_config.latents_mean + self.config = config.arch_config + self.encoder = QwenImageEncoder3d( - base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, input_channels=self.input_channels ) self.quant_conv = QwenImageCausalConv3d(z_dim * 2, z_dim * 2, 1) self.post_quant_conv = QwenImageCausalConv3d(z_dim, z_dim, 1) self.decoder = QwenImageDecoder3d( - base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, input_channels=self.input_channels ) self.spatial_compression_ratio = 2 ** len(self.temperal_downsample) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py index d08c6ed37..2f1dd2913 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py @@ -3,6 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 from diffusers.image_processor import VaeImageProcessor +from sglang.multimodal_gen.runtime.models.model_stages.qwen_image_layered import ( + QwenImageLayeredBeforeDenoisingStage, +) from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, @@ -196,4 +199,62 @@ class QwenImageEditPlusPipeline(QwenImageEditPipeline): pipeline_name = "QwenImageEditPlusPipeline" -EntryClass = [QwenImagePipeline, QwenImageEditPipeline, QwenImageEditPlusPipeline] +def prepare_mu_layered(batch: Req, server_args: ServerArgs): + base_seqlen = 256 * 256 / 16 / 16 + mu = (batch.image_latent.shape[1] / base_seqlen) ** 0.5 + return "mu", mu + + +class QwenImageLayeredPipeline(QwenImageEditPipeline): + pipeline_name = "QwenImageLayeredPipeline" + + _required_config_modules = [ + "vae", + "tokenizer", + "processor", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="QwenImageLayeredBeforeDenoisingStage", + stage=QwenImageLayeredBeforeDenoisingStage( + vae=self.get_module("vae"), + tokenizer=self.get_module("tokenizer"), + processor=self.get_module("processor"), + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + model_path=self.model_path, + ), + ) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu_layered], + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = [ + QwenImagePipeline, + QwenImageEditPipeline, + QwenImageEditPlusPipeline, + QwenImageLayeredPipeline, +]