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

Co-authored-by: chengyusong1 <chengyusong1@jd.com>

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"

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",

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)

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

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

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)

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

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

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 (
+        if (
-            self.enable_sequence_shard is None or self.enable_sequence_shard
+            "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

@@ -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

@@ -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/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

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):
+    def encoding_image_edit(self, outputs, image_inputs, pipeline_config):
-        # encoder hidden state
+        """Encode image-edit text features via pipeline-configured postprocess hook."""
-        prompt_embeds = qwen_image_postprocess_text(outputs, image_inputs, 64)
+        postprocess_funcs = getattr(pipeline_config, "postprocess_text_funcs", ())
-        return prompt_embeds
+        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

@@ -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

@@ -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

@@ -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"