diff --git a/python/sglang/multimodal_gen/configs/models/dits/joy_image.py b/python/sglang/multimodal_gen/configs/models/dits/joy_image.py new file mode 100644 index 000000000..7e9ee6006 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/joy_image.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class JoyImageArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + # Condition embedder mappings + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + # Double blocks mappings + r"^double_blocks\.(\d+)\.attn\.(.*)$": r"double_blocks.\1.\2", + r"^double_blocks\.(\d+)\.img_mlp\.net\.0\.proj\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2", + r"^double_blocks\.(\d+)\.img_mlp\.net\.2\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2", + r"^double_blocks\.(\d+)\.txt_mlp\.net\.0\.proj\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2", + r"^double_blocks\.(\d+)\.txt_mlp\.net\.2\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2", + r"^double_blocks\.(\d+)\.img_attn_qkv\.(.*)$": r"double_blocks.\1.img_attn_qkv.\2", + r"^double_blocks\.(\d+)\.txt_attn_qkv\.(.*)$": r"double_blocks.\1.txt_attn_qkv.\2", + r"^double_blocks\.(\d+)\.img_attn_proj\.(.*)$": r"double_blocks.\1.img_attn_proj.\2", + r"^double_blocks\.(\d+)\.txt_attn_proj\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2", + r"^double_blocks\.(\d+)\.img_mod\.(.*)$": r"double_blocks.\1.img_mod.\2", + r"^double_blocks\.(\d+)\.txt_mod\.(.*)$": r"double_blocks.\1.txt_mod.\2", + r"^double_blocks\.(\d+)\.img_attn_q_norm\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2", + r"^double_blocks\.(\d+)\.img_attn_k_norm\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2", + r"^double_blocks\.(\d+)\.txt_attn_q_norm\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2", + r"^double_blocks\.(\d+)\.txt_attn_k_norm\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2", + } + ) + + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Model architecture parameters + patch_size: tuple[int, int, int] = (1, 2, 2) + num_attention_heads: int = 32 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + mm_double_blocks_depth: int = 40 + freq_dim: int = 256 + text_states_dim: int = 4096 + mlp_width_ratio: float = 4.0 + rope_theta: int = 10000 + rope_dim_list: list[int] = field(default_factory=lambda: [16, 56, 56]) + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class JoyImageDiTConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=JoyImageArchConfig) + prefix: str = "JoyImage" diff --git a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py index 410ba3a02..29b40f957 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py @@ -19,6 +19,7 @@ from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig +from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config __all__ = [ @@ -33,6 +34,7 @@ __all__ = [ "build_flux2_text_messages", "LlamaConfig", "Qwen3TextConfig", + "Qwen3VLConfig", "T5Config", "Gemma2Config", "Gemma3Config", diff --git a/python/sglang/multimodal_gen/configs/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/configs/models/encoders/qwen3vl.py new file mode 100644 index 000000000..d23d269fb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/encoders/qwen3vl.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +def _is_transformer_layer(n: str, m) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m) -> bool: + return n.endswith("norm") + + +@dataclass +class Qwen3VLArchConfig(TextEncoderArchConfig): + """Architecture configuration for Qwen3-VL text encoder. + + Qwen3-VL-8B-Instruct is used by JoyImage model. + Architecture is similar to Qwen2.5-VL but with Qwen3 improvements. + """ + + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int = -1 + eos_token_id: int = 2 + pretraining_tp: int = 1 + tie_word_embeddings: bool = False + rope_theta: float = 10000.0 + rope_scaling: float | None = None + attention_bias: bool = False + attention_dropout: float = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + hidden_state_skip_layer: int = 2 + text_len: int = 2048 + + stacked_params_mapping: list[tuple[str, str, str]] = field( + default_factory=lambda: [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + ) + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] + ) + + # JoyImage specific settings + text_token_max_length: int = 2048 + prompt_template_encode_start_idx = { + "image": 34, + "video": 91, + } + + def __post_init__(self): + super().__post_init__() + self.tokenizer_kwargs = { + "padding": True, + "truncation": True, + "max_length": self.text_len + + self.prompt_template_encode_start_idx["image"], + "return_tensors": "pt", + } + + +@dataclass +class Qwen3VLConfig(TextEncoderConfig): + """Configuration for Qwen3-VL text encoder. + + Used by JoyImage model. + """ + + arch_config: TextEncoderArchConfig = field(default_factory=Qwen3VLArchConfig) diff --git a/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py b/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py index f61f67dc9..caf59c407 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py @@ -89,3 +89,8 @@ class WanVAEConfig(VAEConfig): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) * 2 + + def get_vae_scale_factor(self): + # Wan VAE does not expose block_out_channels like SD-style VAEs. + # Its spatial downsample factor is explicitly defined by scale_factor_spatial. + return self.arch_config.scale_factor_spatial diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 26e007821..da652a7ae 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -876,6 +876,8 @@ class ImagePipelineConfig(PipelineConfig): def shard_latents_for_sp(self, batch, latents): # latents: [B, H * W, C] sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() + if batch.enable_sequence_shard: + return latents, False seq_len = latents.shape[1] # TODO: reuse code in PipelineConfig::shard_latents_for_sp diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/joy_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/joy_image.py new file mode 100644 index 000000000..099ad699b --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/joy_image.py @@ -0,0 +1,431 @@ +import math +from dataclasses import dataclass, field +from typing import Callable, Tuple + +import torch +import torchvision.transforms.functional as TF +from einops import rearrange +from PIL import Image + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.joy_image import JoyImageDiTConfig +from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig +from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, +) + + +def joy_image_postprocess_text( + outputs, + _text_inputs, + drop_idx=34, + max_sequence_length=4096, +): + last_hidden_states = outputs.hidden_states[-1] + prompt_embeds = last_hidden_states[:, drop_idx:] + if max_sequence_length is not None and prompt_embeds.shape[1] > max_sequence_length: + prompt_embeds = prompt_embeds[:, -max_sequence_length:, :] + return prompt_embeds + + +@dataclass +class JoyImageEditPipelineConfig(ImagePipelineConfig): + task_type: ModelTaskType = ModelTaskType.I2I + + dit_config: DiTConfig = field(default_factory=JoyImageDiTConfig) + + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + flow_shift: float = 1.5 + + # Text encoding stage (Qwen3-VL for both text and image understanding) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (Qwen3VLConfig(),) + ) + + enable_torch_compile: bool = False + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) + postprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (joy_image_postprocess_text,) + ) + prioritize_frame_matching: bool = True + bucket_configs: list[tuple[int, int, int, int, int]] = field(init=False) + + def __post_init__(self): + self.bucket_configs = self.generate_video_image_bucket( + basesize=1024, + min_temporal=1, + max_temporal=1, + bs_img=8, + bs_vid=4, + bs_mimg=8, + min_items=1, + max_items=6, + ) + + def slice_noise_pred(self, noise, latents): + # remove noise over input image + noise = noise[:, : latents.size(1)] + return noise + + def _generate_hw_buckets( + self, + base_height=256, + base_width=256, + step_width=16, + step_height=16, + max_ratio=4.0, + ) -> list[tuple[int, int, int, int, int]]: + """Generate dimension buckets based on aspect ratios""" + buckets = [] + target_pixels = base_height * base_width + + height = target_pixels // step_width + width = step_width + + while height >= step_height: + if max(height, width) / min(height, width) <= max_ratio: + ratio = height / width + buckets.append((1, 1, 1, height, width)) + # Try to increase width or decrease height + if height * (width + step_width) <= target_pixels: + width += step_width + else: + height -= step_height + + return buckets + + def generate_video_image_bucket( + self, + basesize=256, + min_temporal=65, + max_temporal=129, + bs_img=8, + bs_vid=1, + bs_mimg=4, + min_items=1, + max_items=1, + ): + # (batch_size, num_items, num_frames, height, width) + assert basesize in [ + 256, + 512, + 768, + 1024, + ], f"[generate_video_image_bucket] wrong basesize {basesize}" + bucket_list = [] + + base_bucket_list = self._generate_hw_buckets() + # image + for _bucket in base_bucket_list: + bucket = list(_bucket) + bucket[0] = bs_img + bucket_list.append(bucket) + # video + for temporal in range(min_temporal, max_temporal + 1, 8): + for _bucket in base_bucket_list: + bucket = list(_bucket) + bs = (max_temporal + 1) // temporal * bs_vid + bucket[0] = bs + bucket[2] = temporal + bucket_list.append(bucket) + # multiple images + for num_items in range(min_items, max_items + 1): + for _bucket in base_bucket_list: + bucket = list(_bucket) + bucket[0] = bs_mimg + bucket[1] = num_items + bucket_list.append(bucket) + # spatial resize + if basesize > 256: + ratio = basesize // 256 + + def resize(bucket, r): + bucket[-2] *= r + bucket[-1] *= r + return bucket + + bucket_list = [resize(bucket, ratio) for bucket in bucket_list] + return bucket_list + + def find_best_bucket( + self, media_shape: tuple[int, int, int, int] + ) -> tuple[int, int, int, int, int]: + """ + Find the best matching bucket for given media dimensions. + + Args: + media_shape: (num_items, num_frames, height, width) of input media + + Returns: + Best matching bucket as (batch_size, num_items, num_frames, height, width) + """ + num_items, num_frames, height, width = media_shape + target_aspect_ratio = height / width + + if num_frames == 1: + valid_buckets = [] + for bucket in self.bucket_configs: + if bucket[1] == num_items and bucket[2] == 1: + valid_buckets.append(bucket) + + if len(valid_buckets) == 0: + raise ValueError(f"No image buckets found for shape {media_shape}") + + return min( + valid_buckets, + key=lambda bucket: abs((bucket[3] / bucket[4]) - target_aspect_ratio), + ) + else: + valid_buckets = [] + for bucket in self.bucket_configs: + if bucket[1] == num_items and bucket[2] > 1 and bucket[2] <= num_frames: + valid_buckets.append(bucket) + + if len(valid_buckets) == 0: + raise ValueError(f"No video buckets found for shape {media_shape}") + + if self.prioritize_frame_matching: + max_frame_count = max(bucket[2] for bucket in valid_buckets) + max_frame_buckets = [ + bucket for bucket in valid_buckets if bucket[2] == max_frame_count + ] + + return min( + max_frame_buckets, + key=lambda bucket: abs( + (bucket[3] / bucket[4]) - target_aspect_ratio + ), + ) + else: + min_ratio_difference = min( + abs((bucket[3] / bucket[4]) - target_aspect_ratio) + for bucket in valid_buckets + ) + best_ratio_buckets = [ + bucket + for bucket in valid_buckets + if abs((bucket[3] / bucket[4]) - target_aspect_ratio) + == min_ratio_difference + ] + + return max(best_ratio_buckets, key=lambda bucket: bucket[2]) + + def resize_center_crop( + self, img: Image.Image, target_size: Tuple[int, int] + ) -> Image.Image: + if isinstance(img, list): + img = img[0] + w, h = img.size # PIL (width, height) + bh, bw = target_size + if w == bw and h == bh: + return img + + scale = max(bh / h, bw / w) + resize_h, resize_w = math.ceil(h * scale), math.ceil(w * scale) + + img = TF.resize( + img, + (resize_h, resize_w), + interpolation=TF.InterpolationMode.BILINEAR, + antialias=True, + ) + img = TF.center_crop(img, target_size) + return img + + def preprocess_condition_image( + self, img, width, height, _vae_image_processor + ) -> None: + target_w, target_h = self.prepare_calculated_size(img) + return self.resize_center_crop(img, (target_h, target_w)), (target_w, target_h) + + def get_decode_scale_and_shift( + self, device, dtype, vae + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Get VAE denormalization scale and shift. + + Args: + device: Target device + dtype: Target dtype + vae: VAE model + + Returns: + Tuple of (scaling_factor, shift_factor) + """ + vae_arch_config = self.vae_config.arch_config + + # Create scale factor: 1.0 / std + scaling_factor = 1.0 / torch.tensor( + vae_arch_config.latents_std, device=device + ).view(1, vae_arch_config.z_dim, 1, 1, 1).to(device, dtype) + + # Create shift factor: mean + shift_factor = ( + torch.tensor(vae_arch_config.latents_mean) + .view(1, vae_arch_config.z_dim, 1, 1, 1) + .to(device, dtype) + ) + + return scaling_factor, shift_factor + + def prepare_calculated_size(self, img: Image.Image) -> Tuple[int, int]: + img_h, img_w = img.size[1], img.size[0] # PIL (w,h) + bucket = self.find_best_bucket((1, 1, img_h, img_w)) + return bucket[-1], bucket[-2] # (width, height) + + def prepare_image_processor_kwargs(self, batch, neg=False) -> dict: + prompt = batch.prompt if not neg else batch.negative_prompt + if prompt is None: + return {} + prompt_list = [prompt] if isinstance(prompt, str) else prompt + image_list = batch.condition_image + if image_list is None: + image_list = [] + elif not isinstance(image_list, list): + image_list = [image_list] + + if len(prompt_list) <= 1: + per_prompt_images = [image_list] + elif len(image_list) <= 1: + per_prompt_images = [list(image_list) for _ in prompt_list] + elif len(image_list) == len(prompt_list): + per_prompt_images = [[image] for image in image_list] + else: + raise ValueError( + "JoyImageEdit expects either one shared condition image or " + "the same number of condition images and prompts." + ) + + prompt_template_encode = ( + "<|im_start|>system\n \\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" + ) + img_prompt_template = "<|vision_start|><|image_pad|><|vision_end|>" + txt = [] + for p, prompt_images in zip(prompt_list, per_prompt_images): + base_img_prompt = img_prompt_template * len(prompt_images) + txt.append(prompt_template_encode.format(base_img_prompt + p)) + return dict(text=txt, padding=True, per_prompt_images=per_prompt_images) + + def prepare_latent_shape(self, batch, batch_size: int, num_frames: int) -> Tuple: + """Prepare latent shape for I2I generation with multi-item support. + + Args: + batch: The request batch + batch_size: Batch size + num_frames: Number of frames (1 for image) + + Returns: + Tuple representing latent shape + """ + + shape = ( + batch_size, + self.vae_config.arch_config.z_dim, # 16 for WanxVAE + 1, + int(batch.height) // self.vae_config.arch_config.scale_factor_spatial, + int(batch.width) // self.vae_config.arch_config.scale_factor_spatial, + ) + + return shape + + def postprocess_image_latent(self, latent_condition, batch): + if latent_condition.dim() == 4: + latent_condition = latent_condition.unsqueeze(0) + elif latent_condition.dim() != 5: + raise ValueError( + f"Expected 4D/5D condition latents, but got shape {latent_condition.shape}" + ) + + batch_size = int(batch.batch_size) + cond_batch = int(latent_condition.shape[0]) + if batch_size > cond_batch: + if batch_size % cond_batch != 0: + raise ValueError( + f"Cannot duplicate condition image latents from batch size {cond_batch} " + f"to target batch size {batch_size}." + ) + repeat_factor = batch_size // cond_batch + latent_condition = latent_condition.repeat(repeat_factor, 1, 1, 1, 1) + elif batch_size < cond_batch: + raise ValueError( + f"Condition image latents batch size {cond_batch} exceeds target batch size {batch_size}." + ) + _, _, t, h, w = latent_condition.shape + pt, ph, pw = self.dit_config.arch_config.patch_size + condition_size = (t // pt, h // ph, w // pw) + + if batch.vae_image_sizes is None: + batch.vae_image_sizes = [condition_size] + else: + # ImageVAEEncodingStage iterates condition images in input order. + # Keep the same order in vae_image_sizes for RoPE range construction. + batch.vae_image_sizes = batch.vae_image_sizes + [condition_size] + + latents = rearrange( + latent_condition, + "b c (t pt) (h ph) (w pw) -> b (t h w) c pt ph pw", + pt=pt, + ph=ph, + pw=pw, + ) + return latents + + def maybe_pack_latents(self, latents, batch_size, batch): + if latents.dim() == 4: + latents = latents.unsqueeze(0) + elif latents.dim() != 5: + raise ValueError(f"Expected 4D/5D latents, but got shape {latents.shape}") + + _, _, t, h, w = latents.shape + pt, ph, pw = self.dit_config.arch_config.patch_size + if batch.vae_image_sizes is None: + batch.vae_image_sizes = [(t // pt, h // ph, w // pw)] + else: + # LatentPreparationStage packs noisy latents after condition latents were packed + # in ImageVAEEncodingStage. Denoising concatenates as [noisy, condition...], + # so keep noisy size at index 0. + batch.vae_image_sizes = [ + (t // pt, h // ph, w // pw) + ] + batch.vae_image_sizes + latents = rearrange( + latents, + "b c (t pt) (h ph) (w pw) -> b (t h w) c pt ph pw", + pt=pt, + ph=ph, + pw=pw, + ) + + return latents + + def post_denoising_loop(self, latents, batch): + lt, lh, lw = batch.vae_image_sizes[0] + target_len = lt * lh * lw + target_patches = latents[:, :target_len] + return rearrange( + target_patches, + "b (t h w) c pt ph pw -> b c (t pt) (h ph) (w pw)", + t=lt, + h=lh, + w=lw, + ) + + def postprocess_cfg_noise( + self, + batch, + noise_pred: torch.Tensor, + noise_pred_cond: torch.Tensor, + ) -> torch.Tensor: + cond_norm = torch.norm(noise_pred_cond, dim=2, keepdim=True) + noise_norm = torch.norm(noise_pred, dim=2, keepdim=True).clamp_min(1e-12) + return noise_pred * (cond_norm / noise_norm) 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 316b7b3a0..4fc40558d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -75,6 +75,10 @@ def qwen_image_postprocess_text( return prompt_embeds +def qwen_image_edit_postprocess_text(outputs, _text_inputs): + return qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=64) + + def _normalize_prompt_list(prompt): return [prompt] if isinstance(prompt, str) else prompt @@ -354,6 +358,9 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): """Configuration for the QwenImageEdit pipeline.""" task_type: ModelTaskType = ModelTaskType.I2I + postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (qwen_image_edit_postprocess_text,) + ) def _prepare_edit_cond_kwargs( self, batch, prompt_embeds, rotary_emb, device, dtype diff --git a/python/sglang/multimodal_gen/configs/sample/joy_image.py b/python/sglang/multimodal_gen/configs/sample/joy_image.py new file mode 100644 index 000000000..857d72b7a --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/joy_image.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class JoyImageEditSamplingParams(SamplingParams): + """Default sampling params for JoyImage Edit single-image I2I.""" + + negative_prompt: str = "" + num_frames: int = 1 + guidance_scale: float = 4.0 + num_inference_steps: int = 40 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 123bf5234..a5dad1fd7 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -493,9 +493,11 @@ class SamplingParams: pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower() - if ("wan" in pipeline_name_lower or "helios" in pipeline_name_lower) and ( - self.enable_sequence_shard is None or self.enable_sequence_shard - ): + if ( + "wan" in pipeline_name_lower + or "helios" in pipeline_name_lower + or "joy" in pipeline_name_lower + ) and (self.enable_sequence_shard is None or self.enable_sequence_shard): self.enable_sequence_shard = True logger.debug("Automatically enabled enable_sequence_shard") else: diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 2b4cb52ea..232384e01 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -55,6 +55,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.glm_image import ( from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.joy_image import ( + JoyImageEditPipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import ( MOVA360PConfig, @@ -97,6 +100,9 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( HunyuanSamplingParams, ) from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams +from sglang.multimodal_gen.configs.sample.joy_image import ( + JoyImageEditSamplingParams, +) from sglang.multimodal_gen.configs.sample.ltx_2 import ( LTX2SamplingParams, LTX23HQSamplingParams, @@ -962,6 +968,18 @@ def _register_configs(): ], ) + # JoyAI + register_configs( + sampling_param_cls=JoyImageEditSamplingParams, + pipeline_config_cls=JoyImageEditPipelineConfig, + hf_model_paths=[ + "jdopensource/JoyAI-Image-Edit-Diffusers", + ], + model_detectors=[ + lambda hf_id: "joyai-image-edit" in hf_id.lower(), + ], + ) + _register_configs() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py b/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py new file mode 100644 index 000000000..545990a5c --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from functools import lru_cache +from typing import Any, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange + +from sglang.multimodal_gen.configs.models.dits.joy_image import JoyImageDiTConfig +from sglang.multimodal_gen.runtime.distributed import ( + get_sp_group, + get_sp_world_size, + sequence_model_parallel_all_gather, +) +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import ( + LayerNormScaleShift, + RMSNorm, + apply_qk_norm_with_optional_rope, +) +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.layers.rotary_embedding import NDRotaryEmbedding +from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.managers.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.models.dits.wanvideo import WanTimeTextImageEmbedding +from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) +_MODULATION_FACTOR = 6 + + +def fused_add_gate( + residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor +) -> torch.Tensor: + """Fused residual addition with gate. + + Computes: residual + x * gate.unsqueeze(1) + + This fuses the gate multiplication and residual addition to reduce + intermediate tensor allocations and memory bandwidth. + + Args: + residual (torch.Tensor): The residual tensor to add to. Shape: (B, L, D) + x (torch.Tensor): The input tensor to be gated. Shape: (B, L, D) + gate (torch.Tensor): The gate tensor. Shape: (B, D) + + Returns: + torch.Tensor: residual + x * gate.unsqueeze(1) + """ + return torch.addcmul(residual, x, gate.unsqueeze(1)) + + +class ModulateWan(nn.Module): + """Modulation layer for WanX.""" + + def __init__(self, hidden_size: int, factor: int, dtype=None, device=None): + super().__init__() + self.factor = factor + self.modulate_table = nn.Parameter( + torch.zeros(1, factor, hidden_size, dtype=dtype, device=device) + / hidden_size**0.5, + requires_grad=False, + ) + set_weight_attrs( + self.modulate_table, + { + "input_dim": 1, + "output_dim": 2, + }, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if len(x.shape) != 3: + x = x.unsqueeze(1) + return [ + o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1) + ] + + +class MMDoubleStreamBlock(nn.Module): + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_pytorch_tanh", + supported_attention_backends: set[AttentionBackendEnum] | None = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.heads_num = heads_num + self.hidden_size = hidden_size + self.head_dim = self.hidden_size // self.heads_num + self.mlp_hidden_dim = int(self.hidden_size * mlp_width_ratio) + + self.img_mod = ModulateWan(self.hidden_size, factor=_MODULATION_FACTOR) + self.fused_modulate_img_norm1 = LayerNormScaleShift( + self.hidden_size, + eps=1e-6, + elementwise_affine=False, + ) + + self.img_attn_qkv = ReplicatedLinear( + self.hidden_size, + hidden_size * 3, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.img_attn_qkv", + ) + self.img_attn_q_norm = RMSNorm( + self.head_dim, + eps=1e-6, + ) + self.img_attn_k_norm = RMSNorm( + self.head_dim, + eps=1e-6, + ) + self.img_attn_proj = ReplicatedLinear( + self.hidden_size, + hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.img_attn_proj", + ) + + self.fused_modulate_img_norm2 = LayerNormScaleShift( + self.hidden_size, + eps=1e-6, + elementwise_affine=False, + ) + self.img_mlp = MLP( + input_dim=self.hidden_size, + mlp_hidden_dim=self.mlp_hidden_dim, + act_type=mlp_act_type, + quant_config=quant_config, + prefix=f"{prefix}.img_mlp", + ) + + # Text modulation and attention + self.txt_mod = ModulateWan(self.hidden_size, factor=_MODULATION_FACTOR) + self.fused_modulate_txt_norm1 = LayerNormScaleShift( + self.hidden_size, + eps=1e-6, + elementwise_affine=False, + ) + self.txt_attn_qkv = ReplicatedLinear( + self.hidden_size, + self.hidden_size * 3, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.txt_attn_qkv", + ) + self.txt_attn_q_norm = RMSNorm( + self.head_dim, + eps=1e-6, + ) + self.txt_attn_k_norm = RMSNorm( + self.head_dim, + eps=1e-6, + ) + self.txt_attn_proj = ReplicatedLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.txt_attn_proj", + ) + + self.fused_modulate_txt_norm2 = LayerNormScaleShift( + self.hidden_size, + eps=1e-6, + elementwise_affine=False, + ) + self.txt_mlp = MLP( + input_dim=self.hidden_size, + mlp_hidden_dim=self.mlp_hidden_dim, + act_type=mlp_act_type, + quant_config=quant_config, + prefix=f"{prefix}.txt_mlp", + ) + self.attn = USPAttention( + num_heads=self.heads_num, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + softmax_scale=None, + ) + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + vis_freqs_cis: Optional[torch.Tensor] = None, + txt_freqs_cis: Optional[torch.Tensor] = None, + num_replicated_suffix: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass through multimodal double stream block.""" + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec) + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec) + + # Image attention + img_modulated = self.fused_modulate_img_norm1( + img, shift=img_mod1_shift, scale=img_mod1_scale + ) + img_qkv, _ = self.img_attn_qkv(img_modulated) + img_q, img_k, img_v = rearrange( + img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num + ) + + if vis_freqs_cis is None: + raise ValueError( + "vis_freqs_cis is required for fused QK-Norm + RoPE kernel" + ) + if not (isinstance(vis_freqs_cis, torch.Tensor) and vis_freqs_cis.dim() == 2): + raise ValueError("vis_freqs_cis must be a 2D cos_sin_cache tensor") + if img_q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError( + f"Fused QK-Norm + RoPE kernel only supports float16/bfloat16, but got {img_q.dtype}" + ) + img_q = img_q.contiguous() + img_k = img_k.contiguous() + img_q, img_k = apply_qk_norm_with_optional_rope( + q=img_q, + k=img_k, + q_norm=self.img_attn_q_norm, + k_norm=self.img_attn_k_norm, + head_dim=img_q.shape[-1], + cos_sin_cache=vis_freqs_cis, + is_neox=False, + allow_inplace=True, + ) + img_q, img_k = img_q.to(img_v), img_k.to(img_v) + + # Text attention + txt_modulated = self.fused_modulate_txt_norm1( + txt, shift=txt_mod1_shift, scale=txt_mod1_scale + ) + txt_qkv, _ = self.txt_attn_qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange( + txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num + ) + + if txt_freqs_cis is not None and not ( + isinstance(txt_freqs_cis, torch.Tensor) and txt_freqs_cis.dim() == 2 + ): + raise ValueError("txt_freqs_cis must be a 2D cos_sin_cache tensor") + txt_q = txt_q.contiguous() + txt_k = txt_k.contiguous() + txt_q, txt_k = apply_qk_norm_with_optional_rope( + q=txt_q, + k=txt_k, + q_norm=self.txt_attn_q_norm, + k_norm=self.txt_attn_k_norm, + head_dim=txt_q.shape[-1], + cos_sin_cache=txt_freqs_cis, + is_neox=False, + allow_inplace=True, + ) + txt_q, txt_k = txt_q.to(txt_v), txt_k.to(txt_v) + + # Attention + joint_query = torch.cat([img_q, txt_q], dim=1) + joint_key = torch.cat([img_k, txt_k], dim=1) + joint_value = torch.cat([img_v, txt_v], dim=1) + attn = self.attn( + joint_query, + joint_key, + joint_value, + num_replicated_suffix=num_replicated_suffix, + ) + attn = attn.flatten(2, 3) + img_attn, txt_attn = ( + attn[:, : img.shape[1]], + attn[:, img.shape[1] :], + ) + + img = fused_add_gate(img, self.img_attn_proj(img_attn)[0], img_mod1_gate) + img = fused_add_gate( + img, + self.img_mlp( + self.fused_modulate_img_norm2( + img, shift=img_mod2_shift, scale=img_mod2_scale + ) + ), + img_mod2_gate, + ) + + # Text blocks + txt = fused_add_gate(txt, self.txt_attn_proj(txt_attn)[0], txt_mod1_gate) + txt = fused_add_gate( + txt, + self.txt_mlp( + self.fused_modulate_txt_norm2( + txt, shift=txt_mod2_shift, scale=txt_mod2_scale + ) + ), + txt_mod2_gate, + ) + + return img, txt + + +class JoyTransformer3DModel(CachableDiT, OffloadableDiTMixin): + """ + JoyImage Transformer 3D Model for image generation. + + """ + + _supports_gradient_checkpointing = True + _fsdp_shard_conditions = JoyImageDiTConfig()._fsdp_shard_conditions + _compile_conditions = JoyImageDiTConfig()._compile_conditions + _supported_attention_backends = JoyImageDiTConfig()._supported_attention_backends + param_names_mapping = JoyImageDiTConfig().param_names_mapping + reverse_param_names_mapping = JoyImageDiTConfig().reverse_param_names_mapping + lora_param_names_mapping = JoyImageDiTConfig().lora_param_names_mapping + + def __init__( + self, + config: JoyImageDiTConfig, + hf_config: dict[str, Any], + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__( + config=config, + hf_config=hf_config, + ) + self.in_channels = config.in_channels + self.out_channels = config.out_channels or config.in_channels + self.patch_size = config.patch_size + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.rope_dim_list = config.rope_dim_list + self.mm_double_blocks_depth = config.mm_double_blocks_depth + self.rope_theta = config.rope_theta + self.quant_config = quant_config + self.num_channels_latents = self.out_channels + + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"Hidden size {self.hidden_size} must be divisible by num_attention_heads {self.num_attention_heads}" + ) + + # Image projection (patch embedding) + self.img_in = nn.Conv3d( + self.in_channels, + self.hidden_size, + kernel_size=self.patch_size, + stride=self.patch_size, + ) + + # Condition embedding + self.condition_embedder = WanTimeTextImageEmbedding( + dim=self.hidden_size, + time_freq_dim=config.freq_dim, + text_embed_dim=config.text_states_dim, + ) + + # Double blocks (DiT layers) + self.double_blocks = nn.ModuleList( + [ + MMDoubleStreamBlock( + self.hidden_size, + self.num_attention_heads, + mlp_width_ratio=config.mlp_width_ratio, + supported_attention_backends=self._supported_attention_backends, + quant_config=quant_config, + prefix=f"{config.prefix}.double_blocks.{i}", + ) + for i in range(self.mm_double_blocks_depth) + ] + ) + # Layerwise offload expects ModuleList names here. + self.layer_names = ["double_blocks"] + + # Output norm & projection + self.norm_out = nn.LayerNorm( + self.hidden_size, elementwise_affine=False, eps=1e-6 + ) + self.proj_out = ReplicatedLinear( + self.hidden_size, + self.out_channels * math.prod(self.patch_size), + quant_config=quant_config, + prefix=f"proj_out", + ) + self.__post_init__() + + self.sp_size = get_sp_world_size() + self.rotary_emb = NDRotaryEmbedding( + rope_dim_list=config.rope_dim_list, + rope_theta=config.rope_theta, + dtype=torch.float32, + ) + + @lru_cache(maxsize=1) + def _compute_rope_for_local_shard( + self, + local_len: int, + rank: int, + vae_image_sizes: tuple[tuple[int, int, int], ...], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + token_start = rank * local_len + token_indices = torch.arange( + token_start, + token_start + local_len, + device=device, + dtype=torch.long, + ) + positions = torch.zeros(local_len, 3, device=device, dtype=torch.long) + + cumsum = 0 + current_t_offset = 0 + for t, h, w in vae_image_sizes: + item_size = t * h * w + mask = (token_indices >= cumsum) & (token_indices < cumsum + item_size) + if mask.any(): + local_idx = token_indices[mask] - cumsum + frame_stride = h * w + positions[mask, 0] = local_idx // frame_stride + current_t_offset + positions[mask, 1] = (local_idx % frame_stride) // w + positions[mask, 2] = local_idx % w + cumsum += item_size + current_t_offset += t + + return self.rotary_emb.forward_uncached(positions) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_mask: torch.Tensor | list[torch.Tensor] | None = None, + vis_freqs_cis: torch.Tensor | None = None, + txt_freqs_cis: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass through JoyImage Transformer.""" + forward_batch = get_forward_context().forward_batch + sequence_shard_enabled = ( + forward_batch is not None + and getattr(forward_batch, "enable_sequence_shard", False) + and self.sp_size > 1 + ) + + batch_size = hidden_states.shape[0] + + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + + if isinstance(encoder_hidden_states_mask, list): + encoder_hidden_states_mask = encoder_hidden_states_mask[0] + + cond_batch = int(encoder_hidden_states.shape[0]) + if cond_batch != int(batch_size): + if cond_batch <= 0 or int(batch_size) % cond_batch != 0: + raise ValueError( + "JoyImage conditioning batch mismatch: " + f"hidden_states batch={batch_size}, " + f"encoder_hidden_states batch={cond_batch}." + ) + repeat_factor = int(batch_size) // cond_batch + encoder_hidden_states = encoder_hidden_states.repeat_interleave( + repeat_factor, dim=0 + ) + if encoder_hidden_states_mask is not None: + encoder_hidden_states_mask = ( + encoder_hidden_states_mask.repeat_interleave(repeat_factor, dim=0) + ) + + # Prepare img + x = rearrange(hidden_states, "b n c p1 p2 p3 -> (b n) c p1 p2 p3") + x = self.img_in(x) + img = rearrange(x, "(b n) d 1 1 1 -> b n d", b=batch_size) + + seq_len_orig = img.shape[1] + seq_shard_pad = 0 + if sequence_shard_enabled: + if seq_len_orig % self.sp_size != 0: + seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) + pad = torch.zeros( + (batch_size, seq_shard_pad, img.shape[2]), + dtype=img.dtype, + device=img.device, + ) + img = torch.cat([img, pad], dim=1) + sp_rank = get_sp_group().rank_in_group + local_seq_len = img.shape[1] // self.sp_size + img = img.view(batch_size, self.sp_size, local_seq_len, img.shape[2])[ + :, sp_rank, :, : + ].contiguous() + + # Compute rope in model for all SP modes + if forward_batch is not None and forward_batch.vae_image_sizes is not None: + vae_image_sizes = tuple(tuple(s) for s in forward_batch.vae_image_sizes) + local_len = img.shape[1] + rank = get_sp_group().rank_in_group if self.sp_size > 1 else 0 + freqs_cos, freqs_sin = self._compute_rope_for_local_shard( + local_len, + rank, + vae_image_sizes, + img.device, + ) + vis_freqs_cis = torch.cat( + [ + freqs_cos.to(dtype=torch.float32).contiguous(), + freqs_sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) + + _, vec, txt, _ = self.condition_embedder(timestep, encoder_hidden_states) + if vec.shape[-1] > self.hidden_size: + vec = vec.unflatten(1, (_MODULATION_FACTOR, -1)) + + txt_suffix_len = txt.shape[1] if sequence_shard_enabled else 0 + + # Pass through DiT blocks + for block in self.double_blocks: + img, txt = block( + img, + txt, + vec, + vis_freqs_cis, + txt_freqs_cis, + num_replicated_suffix=txt_suffix_len, + ) + + if sequence_shard_enabled: + img = img.contiguous() + img = sequence_model_parallel_all_gather(img, dim=1) + if seq_shard_pad > 0: + img = img[:, :seq_len_orig, :] + + img, _ = self.proj_out(self.norm_out(img)) + + # Restore patch layout expected by downstream latent decoding. + img = rearrange( + img, + "b n (pt ph pw c) -> b n c pt ph pw", + pt=self.patch_size[0], + ph=self.patch_size[1], + pw=self.patch_size[2], + c=self.out_channels, + ) + + return img + + +class JoyImageEditTransformer3DModel(JoyTransformer3DModel): + """Backward-compatible alias for JoyImageEdit model configs.""" + + pass + + +EntryClass = [JoyTransformer3DModel, JoyImageEditTransformer3DModel] diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py new file mode 100644 index 000000000..d1cab5ae9 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py @@ -0,0 +1,1005 @@ +# SPDX-License-Identifier: Apache-2.0 + +from transformers import ( + Cache, + DynamicCache, +) +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.utils import TransformersKwargs, is_torchdynamo_compiling + +from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention +from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader +from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +"""Inference-only Qwen3-VL model compatible with HuggingFace weights.""" +import logging +from typing import Iterable, Optional, Tuple, Union + +try: + from typing import Unpack # type: ignore[attr-defined] +except ImportError: + # Python 3.10 and below + from typing_extensions import Unpack + +import torch +import torch.nn as nn +from transformers.activations import ACT2FN + +logger = logging.getLogger(__name__) + +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLTextConfig, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLCausalLMOutputWithPast, + Qwen3VLModelOutputWithPast, + Qwen3VLTextRMSNorm, + Qwen3VLTextRotaryEmbedding, + Qwen3VLVisionModel, + apply_rotary_pos_emb, +) + + +class Qwen3VLTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = config.hidden_size // config.num_attention_heads + self.num_key_value_groups = ( + config.num_attention_heads // config.num_key_value_heads + ) + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = Qwen3VLTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # unlike olmo, only on the head dim! + self.k_norm = Qwen3VLTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # thus post q_norm does not need reshape + + self.attn = LocalAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_key_value_heads, + softmax_scale=self.scaling, + causal=True, + supported_attention_backends=( + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm( + self.q_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + key_states = self.k_norm( + self.k_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # attention_interface: Callable = eager_attention_forward + """ + _global_mapping = { + "flash_attention_3": flash_attention_forward, + "flash_attention_2": flash_attention_forward, + "flex_attention": flex_attention_forward, + "paged_attention": paged_attention_forward, + "sdpa": sdpa_attention_forward, + "sdpa_paged": sdpa_attention_paged_forward, + "eager_paged": eager_paged_attention_forward, + } + """ + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + attn_output = self.attn(query_states, key_states, value_states) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output + + +class Qwen3VLTextMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Qwen3VLTextDecoderLayer(nn.Module): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx) + + self.mlp = Qwen3VLTextMLP(config) + self.input_layernorm = Qwen3VLTextRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3VLTextRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class Qwen3VLTextModel(nn.Module): + config: Qwen3VLTextConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer"] + + def __init__(self, config: Qwen3VLTextConfig): + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + Qwen3VLTextDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[tuple, BaseModelOutputWithPast]: + r""" + visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*): + The mask of the visual positions. + deepstack_visual_embeds (`list[torch.Tensor]`, *optional*): + The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). + The feature is extracted from the different visual encoder layers, and fed to the decoder + hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334). + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + # the hard-coded `3` is for temporal, height and width. + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1 + ) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + position_ids = position_ids[1:] + else: + text_position_ids = position_ids[0] + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=text_position_ids, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + # decoder layers + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + cache_position=cache_position, + output_attentions=output_attentions, + position_embeddings=position_embeddings, + **kwargs, + ) + # hidden_states = layer_outputs + + # add visual features to the hidden states of first several layers + if deepstack_visual_embeds is not None and layer_idx in range( + len(deepstack_visual_embeds) + ): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states = self.norm(hidden_states) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + past_key_values, + all_hidden_states, + all_self_attns, + ] + if v is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _deepstack_process( + self, + hidden_states: torch.Tensor, + visual_pos_masks: torch.Tensor, + visual_embeds: torch.Tensor, + ): + visual_pos_masks = visual_pos_masks.to(hidden_states.device) + visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) + local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds + hidden_states[visual_pos_masks, :] = local_this + return hidden_states + + +class Qwen3VLModel(nn.Module): + base_model_prefix = "" + _checkpoint_conversion_mapping = {} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + + def __init__(self, config): + super().__init__() + self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.language_model = Qwen3VLTextModel(config.text_config) + self.rope_deltas = None # cache rope_deltas here + self.config = config + + # Initialize weights and apply final processing + + def get_input_embeddings(self): + return self.language_model.embed_tokens + + def set_input_embeddings(self, value): + self.language_model.embed_tokens = value + + def set_decoder(self, decoder): + self.language_model = decoder + + def get_decoder(self): + return self.language_model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids.""" + + # Since we use timestamps to separate videos, like , the video_grid_thw should also be split + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0], dim=0 + ) + video_grid_thw[:, 0] = 1 + + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: Optional[torch.LongTensor] = None, + ): + """ + Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + + Args: + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + # Same implementation as for images + return self.get_image_features(pixel_values_videos, video_grid_thw) + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + ): + """ + Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + pixel_values = pixel_values.type(self.visual.dtype) + visual_out = self.visual(pixel_values, grid_thw=image_grid_thw) + image_embeds = visual_out.pooler_output + deepstack_image_embeds = visual_out.deepstack_features + split_sizes = ( + image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2 + ).tolist() + image_embeds = torch.split(image_embeds, split_sizes) + return image_embeds, deepstack_image_embeds + + def get_placeholder_mask( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_features: Optional[torch.FloatTensor] = None, + video_features: Optional[torch.FloatTensor] = None, + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor( + self.config.image_token_id, + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + special_image_mask = special_image_mask.all(-1) + special_video_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor( + self.config.video_token_id, + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + special_video_mask = special_video_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + special_video_mask = input_ids == self.config.video_token_id + + n_image_tokens = special_image_mask.sum() + special_image_mask = ( + special_image_mask.unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + if ( + image_features is not None + and inputs_embeds[special_image_mask].numel() != image_features.numel() + ): + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" + ) + + n_video_tokens = special_video_mask.sum() + special_video_mask = ( + special_video_mask.unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + if ( + video_features is not None + and inputs_embeds[special_video_mask].numel() != video_features.numel() + ): + raise ValueError( + f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}" + ) + + return special_image_mask, special_video_mask + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + use_cache: Optional[bool] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Union[tuple, Qwen3VLModelOutputWithPast]: + r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + image_mask = None + video_mask = None + + if pixel_values is not None: + image_embeds, deepstack_image_embeds = self.get_image_features( # long + pixel_values, image_grid_thw + ) + image_embeds = torch.cat(image_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_embeds, deepstack_video_embeds = self.get_video_features( + pixel_values_videos, video_grid_thw + ) + video_embeds = torch.cat(video_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + visual_pos_masks = None + deepstack_visual_embeds = None + if image_mask is not None and video_mask is not None: + # aggregate visual_pos_masks and deepstack_visual_embeds + image_mask = image_mask[..., 0] + video_mask = video_mask[..., 0] + visual_pos_masks = image_mask | video_mask + deepstack_visual_embeds = [] + image_mask_joint = image_mask[visual_pos_masks] + video_mask_joint = video_mask[visual_pos_masks] + for img_embed, vid_embed in zip( + deepstack_image_embeds, deepstack_video_embeds + ): + embed_joint = img_embed.new_zeros( + visual_pos_masks.sum(), img_embed.shape[-1] + ).to(img_embed.device) + embed_joint[image_mask_joint, :] = img_embed + embed_joint[video_mask_joint, :] = vid_embed + deepstack_visual_embeds.append(embed_joint) + elif image_mask is not None: + image_mask = image_mask[..., 0] + visual_pos_masks = image_mask + deepstack_visual_embeds = deepstack_image_embeds + elif video_mask is not None: + video_mask = video_mask[..., 0] + visual_pos_masks = video_mask + deepstack_visual_embeds = deepstack_video_embeds + + if position_ids is None: + attention_mask_tensor = ( + attention_mask + if not isinstance(attention_mask, dict) + else attention_mask["full_attention"] + ) + if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4: + attention_mask_tensor = torch.diagonal( + attention_mask_tensor[:, 0], dim1=1, dim2=2 + ) + # Only apply conversion for floating point tensors (inverted masks) + if attention_mask_tensor.dtype.is_floating_point: + attention_mask_tensor = ( + attention_mask_tensor + / torch.finfo(attention_mask_tensor.dtype).min + ) + attention_mask_tensor = (1.0 - attention_mask_tensor).int() + + # Calculate RoPE index once per generation in the pre-fill stage only. + # When compiling, we can't check tensor values thus we check only input length + # It is safe to assume that `length!=1` means we're in pre-fill because compiled + # models currently cannot do assisted decoding + prefill_compiled_stage = is_torchdynamo_compiling() and ( + (input_ids is not None and input_ids.shape[1] != 1) + or (inputs_embeds is not None and inputs_embeds.shape[1] != 1) + ) + prefill_noncompiled_stage = not is_torchdynamo_compiling() and ( + (cache_position is not None and cache_position[0] == 0) + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ) + if ( + prefill_compiled_stage or prefill_noncompiled_stage + ) or self.rope_deltas is None: + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + attention_mask=attention_mask_tensor, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + **kwargs, + ) + + output = Qwen3VLModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + return output if return_dict else output.to_tuple() + + +class Qwen3VLForConditionalGeneration(TextEncoder): + default_bitsandbytes_target_modules = [ + ".gate_up_proj.", + ".down_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + _checkpoint_conversion_mapping = {} + _tied_weights_keys = ["lm_head.weight"] + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLConfig + + def __init__(self, config): + super().__init__(config) + config = config.arch_config + self.model = Qwen3VLModel(config) + self.lm_head = nn.Linear( + config.text_config.hidden_size, config.text_config.vocab_size, bias=False + ) + + @torch.no_grad() + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + use_cache: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Union[tuple, Qwen3VLCausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + output_attentions = False + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs[0] + + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = ( + slice(-logits_to_keep, None) + if isinstance(logits_to_keep, int) + else logits_to_keep + ) + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + return Qwen3VLCausalLMOutputWithPast( + loss=None, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=outputs.rope_deltas, + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + loaded_params: set[str] = set() + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + try: + param = params_dict[name] + except KeyError: + raise KeyError( + f"Unexpected weight name while loading Qwen3VL: {name}" + ) from None + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight.to(param.dtype) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +EntryClass = Qwen3VLForConditionalGeneration diff --git a/python/sglang/multimodal_gen/runtime/pipelines/joy_image.py b/python/sglang/multimodal_gen/runtime/pipelines/joy_image.py new file mode 100644 index 000000000..76a7307bf --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/joy_image.py @@ -0,0 +1,30 @@ +from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +class JoyImageEditPipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "JoyImageEditPipeline" + + _required_config_modules = [ + "processor", + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + + self.add_standard_ti2i_stages( + vae_image_processor=None, + prompt_encoding="image_encoding", + image_processor_key="processor", + prompt_text_encoder_key="text_encoder", + ) + + +EntryClass = JoyImageEditPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index fe6187c0b..7a7de8c4b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -18,9 +18,6 @@ import torch from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput -from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( - qwen_image_postprocess_text, -) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.component_manager import ComponentUse from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context @@ -149,10 +146,15 @@ class ImageEncodingStage(PipelineStage): uses.append(ComponentUse(stage_name, "text_encoder")) return uses - def encoding_qwen_image_edit(self, outputs, image_inputs): - # encoder hidden state - prompt_embeds = qwen_image_postprocess_text(outputs, image_inputs, 64) - return prompt_embeds + def encoding_image_edit(self, outputs, image_inputs, pipeline_config): + """Encode image-edit text features via pipeline-configured postprocess hook.""" + postprocess_funcs = getattr(pipeline_config, "postprocess_text_funcs", ()) + if not postprocess_funcs or not callable(postprocess_funcs[0]): + raise ValueError( + "Image-edit pipeline requires a callable postprocess_text_funcs[0]." + ) + + return postprocess_funcs[0](outputs, image_inputs) @torch.no_grad() def forward( @@ -262,11 +264,15 @@ class ImageEncodingStage(PipelineStage): ) all_prompt_embeds.append( - self.encoding_qwen_image_edit(outputs, image_inputs) + self.encoding_image_edit( + outputs, image_inputs, server_args.pipeline_config + ) ) if batch.do_classifier_free_guidance: all_neg_prompt_embeds.append( - self.encoding_qwen_image_edit(neg_outputs, neg_image_inputs) + self.encoding_image_edit( + neg_outputs, neg_image_inputs, server_args.pipeline_config + ) ) if all_prompt_embeds: diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 93ad2f55f..31332ca04 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -28,6 +28,7 @@ from sglang.multimodal_gen.test.test_utils import ( DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, + DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST, DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST, @@ -155,6 +156,12 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [ ), MULTI_FRAME_I2I_sampling_params, ), + DiffusionTestCase( + "joyai_image_edit_ti2i", + DiffusionServerArgs(model_path=DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST), + TI2I_sampling_params, + run_consistency_check=False, + ), # Upscaling (Real-ESRGAN 4×) for T2I DiffusionTestCase( "flux_2_image_t2i_upscaling_4x", diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index dacfe4508..12a74783d 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -753,6 +753,62 @@ "expected_median_denoise_ms": 647.87, "estimated_full_test_time_s": 153.6 }, + "joyai_image_edit_ti2i": { + "stages_ms": { + "InputValidationStage": 32.2, + "ImageEncodingStage": 948.69, + "ImageVAEEncodingStage": 70.47, + "LatentPreparationStage": 0.17, + "TimestepPreparationStage": 20.66, + "DenoisingStage": 26894.18, + "DecodingStage": 14.27 + }, + "denoise_step_ms": { + "0": 432.34, + "1": 673.28, + "2": 658.38, + "3": 677.9, + "4": 677.87, + "5": 665.09, + "6": 680.25, + "7": 678.54, + "8": 675.02, + "9": 683.38, + "10": 679.11, + "11": 674.55, + "12": 681.24, + "13": 680.79, + "14": 678.81, + "15": 680.94, + "16": 680.89, + "17": 678.07, + "18": 679.9, + "19": 682.67, + "20": 678.41, + "21": 679.92, + "22": 681.07, + "23": 679.93, + "24": 682.35, + "25": 680.8, + "26": 681.19, + "27": 682.05, + "28": 681.34, + "29": 680.8, + "30": 675.57, + "31": 679.21, + "32": 679.67, + "33": 675.05, + "34": 681.63, + "35": 678.62, + "36": 675.68, + "37": 678.58, + "38": 679.01, + "39": 677.64 + }, + "expected_e2e_ms": 28350.2, + "expected_avg_denoise_ms": 672.19, + "expected_median_denoise_ms": 679.16 + }, "qwen_image_t2i_cache_dit_enabled": { "stages_ms": { "InputValidationStage": 0.06, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 71c96f71b..3397a0a80 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -75,6 +75,9 @@ DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2509" DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2511" DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Layered" +# JoyAI image editing models +DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST = "jdopensource/JoyAI-Image-Edit-Diffusers" + # FLUX image generation models DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.1-dev" DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.2-dev"