minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] model: support JoyAI-Image-Edit (#22625)

Browse Source 
Co-authored-by: chengyusong1 <chengyusong1@jd.com>
This commit is contained in:
Clay
2026-05-02 14:08:57 +08:00
committed by GitHub
co-authored by chengyusong1
parent cd27baaffd
commit 5ec3b26799
17 changed files with 2344 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/joy_image.py
+68
View File
@@ -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"
python/sglang/multimodal_gen/configs/models/encoders/__init__.py
+2
View File
@@ -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",
python/sglang/multimodal_gen/configs/models/encoders/qwen3vl.py
+94
View File
@@ -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)
python/sglang/multimodal_gen/configs/models/vaes/wanvae.py
+5
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+2
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/joy_image.py
+431
View File
@@ -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)
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
+7
View File
@@ -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
python/sglang/multimodal_gen/configs/sample/joy_image.py
+13
View File
@@ -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
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+5 -3
View File
@@ -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:
python/sglang/multimodal_gen/registry.py
+18
View File
@@ -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()
python/sglang/multimodal_gen/runtime/models/dits/joy_image.py
+583
View File
@@ -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]
python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py
+1005
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/pipelines/joy_image.py
+30
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
+15 -9
View File
@@ -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:
python/sglang/multimodal_gen/test/server/gpu_cases.py
+7
View File
@@ -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",
python/sglang/multimodal_gen/test/server/perf_baselines.json
+56
View File
@@ -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,
python/sglang/multimodal_gen/test/test_utils.py
+3
View File
@@ -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"