[diffusion] feat: support LongCat-Image-Edit and LongCat-Image-Edit-Turbo (#35829)

Co-authored-by: 登辉 <yangdenghui.ydh@alibaba-inc.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

docs/docs/sglang-diffusion/compatibility_matrix.mdx

@@ -39,7 +39,7 @@ Rows are grouped when a family shares the same runtime path or optimization supp
           </tr>
           <tr>
             <td>LongCat-Image</td>
-            <td><div className="sgd-id-list"><code>meituan-longcat/LongCat-Image</code></div></td>
+            <td><div className="sgd-id-list"><code>meituan-longcat/LongCat-Image</code><code>meituan-longcat/LongCat-Image-Edit</code><code>meituan-longcat/LongCat-Image-Edit-Turbo</code></div></td>
           </tr>
           <tr>
             <td>SD3 / SD3.5</td>

python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py

@@ -1,3 +1,4 @@
+import math
 from dataclasses import dataclass, field
 from typing import Callable
 
@@ -20,6 +21,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
 from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
     ModelDeploymentConfig,
 )
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_world_size,
+    model_parallel_is_initialized,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -252,6 +257,28 @@ class LongCatImagePipelineConfig(ImagePipelineConfig):
             keep_resident_components=("text_encoder", "vae"),
         )
 
+    def expand_conditioning_to_sample_batch(self, batch):
+        # Noise/reference latents are built at batch_size = prompts * num_outputs,
+        # but text encoding stays per-prompt; repeat the embeds to match. No-op
+        # for num_outputs == 1. Shared by T2I and Edit.
+        from sglang.multimodal_gen.runtime.utils.condition_expansion import (
+            PromptToSampleBatchExpander,
+        )
+
+        expander = PromptToSampleBatchExpander.from_batch(batch)
+        if expander is None:
+            return batch
+        for field_name in (
+            "prompt_embeds",
+            "negative_prompt_embeds",
+            "prompt_embeds_mask",
+            "negative_prompt_embeds_mask",
+            "prompt_seq_lens",
+            "negative_prompt_seq_lens",
+        ):
+            expander.expand_field(batch, field_name)
+        return batch
+
     # --- LatentPreparationStage hooks ---
 
     def prepare_latent_shape(self, batch, batch_size, num_frames):
@@ -416,3 +443,164 @@ class LongCatImagePipelineConfig(ImagePipelineConfig):
         noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
         scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0)
         return noise_pred * scale
+
+
+def _calculate_edit_dimensions(target_area, ratio):
+    """Output size for LongCat-Image-Edit: fit `target_area`, ceil to /16.
+
+    Copied from diffusers pipeline_longcat_image_edit.calculate_dimensions.
+    Note this intentionally differs from sglang.multimodal_gen.utils
+    calculate_dimensions (which rounds to /32).
+    """
+    width = math.sqrt(target_area * ratio)
+    height = width / ratio
+
+    width = width if width % 16 == 0 else (width // 16 + 1) * 16
+    height = height if height % 16 == 0 else (height // 16 + 1) * 16
+
+    return int(width), int(height)
+
+
+@dataclass
+class LongCatImageEditPipelineConfig(LongCatImagePipelineConfig):
+    """Configuration for the LongCat-Image-Edit I2I pipeline.
+
+    Mirrors diffusers LongCatImageEditPipeline: the reference image is resized
+    to the output resolution, VAE-encoded with argmax sampling, packed, and
+    concatenated after the noisy latents along the sequence dim. RoPE ids use
+    modality 1 for noisy tokens and modality 2 for reference tokens, both
+    offset by the full text sequence length (VL image tokens + 512 body).
+    """
+
+    task_type: ModelTaskType = ModelTaskType.I2I
+    # diffusers reference draws latent noise with a CPU generator
+    generator_device: str = "cpu"
+
+    # --- InputValidationStage hooks ---
+
+    def calculate_condition_image_size(self, image, width, height):
+        return _calculate_edit_dimensions(1024 * 1024, width / height)
+
+    # --- LatentPreparationStage hooks ---
+
+    def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype:
+        # The edit reference draws noise directly in the prompt-embeds dtype
+        # (bf16), unlike the T2I reference which draws in float32.
+        return prompt_dtype
+
+    def maybe_prepare_latent_ids(self, latents):
+        # img_ids (noisy + reference) are built per-step in
+        # prepare_*_cond_kwargs because the text start offset depends on the
+        # VL image token count, which is unknown at latent preparation time.
+        return None
+
+    # --- ImageVAEEncodingStage hooks ---
+
+    def preprocess_vae_encode(self, image, vae):
+        # AutoencoderKL is a 2D image VAE; drop the frames dim added by
+        # ImageVAEEncodingStage ([B, C, 1, H, W] -> [B, C, H, W]).
+        if image.dim() == 5 and image.shape[2] == 1:
+            image = image.squeeze(2)
+        return image
+
+    def postprocess_image_latent(self, latent_condition, batch):
+        if latent_condition.dim() == 5 and latent_condition.shape[2] == 1:
+            latent_condition = latent_condition.squeeze(2)
+        batch_size = batch.batch_size
+        if batch_size > latent_condition.shape[0]:
+            if batch_size % latent_condition.shape[0] != 0:
+                raise ValueError(
+                    f"Cannot duplicate reference image of batch size "
+                    f"{latent_condition.shape[0]} to {batch_size} prompts."
+                )
+            latent_condition = latent_condition.repeat(
+                batch_size // latent_condition.shape[0], 1, 1, 1
+            )
+        _, num_channels_latents, height, width = latent_condition.shape
+        return _pack_latents(
+            latent_condition, batch_size, num_channels_latents, height, width
+        )
+
+    # --- Denoising hooks ---
+
+    def shard_latents_for_sp(self, batch, latents):
+        # (h/2)*(w/2) is odd at most ~1MP edit resolutions, so SP has to pad, and
+        # the pads stay unmasked (USPAttention rejects a mask alongside the
+        # replicated text prefix). Repeat the last token instead of the base
+        # class's zeros, which would carry the RoPE of text token 0.
+        if latents.dim() == 3 and model_parallel_is_initialized():
+            sp_world_size = get_sp_world_size()
+            remainder = latents.shape[1] % sp_world_size
+            if remainder:
+                pad = latents[:, -1:].expand(-1, sp_world_size - remainder, -1)
+                latents = torch.cat([latents, pad], dim=1)
+        return super().shard_latents_for_sp(batch, latents)
+
+    def _maybe_shard_pos_ids_for_sp(self, batch, pos_ids):
+        # RoPE ids must be sharded exactly like their latents, so reuse the same
+        # helper. It reads the SP group, which only exists under model parallelism.
+        if not model_parallel_is_initialized() or get_sp_world_size() == 1:
+            return pos_ids
+        sharded, _ = self.shard_latents_for_sp(batch, pos_ids.unsqueeze(0))
+        return sharded.squeeze(0)
+
+    def _edit_img_ids(self, batch, num_token, device):
+        """Position ids for [noisy | reference] packed latent tokens.
+
+        The reference image is resized to the output resolution, so both grids
+        share the same shape; they differ only in modality id (1 vs 2).
+        """
+        vae_scale_factor = self.vae_config.get_vae_scale_factor()
+        h = 2 * (int(batch.height) // (vae_scale_factor * 2))
+        w = 2 * (int(batch.width) // (vae_scale_factor * 2))
+        noisy_ids = _prepare_pos_ids(
+            modality_id=1,
+            token_type="image",
+            start=(num_token, num_token),
+            height=h // 2,
+            width=w // 2,
+        )
+        ref_ids = _prepare_pos_ids(
+            modality_id=2,
+            token_type="image",
+            start=(num_token, num_token),
+            height=h // 2,
+            width=w // 2,
+        )
+        noisy_ids = self._maybe_shard_pos_ids_for_sp(batch, noisy_ids)
+        ref_ids = self._maybe_shard_pos_ids_for_sp(batch, ref_ids)
+        return torch.cat([noisy_ids, ref_ids], dim=0).to(device)
+
+    def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        num_token = batch.prompt_embeds[0].shape[1]
+        return {
+            "txt_ids": _prepare_pos_ids(
+                modality_id=0, token_type="text", start=(0, 0), num_token=num_token
+            ).to(device),
+            "img_ids": self._edit_img_ids(batch, num_token, device),
+        }
+
+    def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        num_token = batch.negative_prompt_embeds[0].shape[1]
+        return {
+            "txt_ids": _prepare_pos_ids(
+                modality_id=0, token_type="text", start=(0, 0), num_token=num_token
+            ).to(device),
+            "img_ids": self._edit_img_ids(batch, num_token, device),
+        }
+
+    def slice_noise_pred(self, noise, latents):
+        # Drop predictions over the appended reference-image tokens.
+        return noise[:, : latents.size(1)]
+
+    def post_denoising_loop(self, latents, batch):
+        # The SP gather leaves the noisy latents at their padded length; trim the
+        # trailing pad tokens before unpacking to the (h/2)*(w/2) grid.
+        if latents.dim() == 3:
+            vae_scale_factor = self.vae_config.get_vae_scale_factor()
+            h = 2 * (int(batch.height) // (vae_scale_factor * 2))
+            w = 2 * (int(batch.width) // (vae_scale_factor * 2))
+            expected = (h // 2) * (w // 2)
+            if latents.shape[1] > expected:
+                latents = latents[:, :expected, :]
+        return super().post_denoising_loop(latents, batch)

python/sglang/multimodal_gen/configs/sample/longcat_image.py

@@ -12,3 +12,28 @@ class LongCatImageSamplingParams(SamplingParams):
     # Override base class defaults to enable LongCat-specific features by default
     enable_cfg_renorm: bool = True
     enable_prompt_rewrite: bool = True
+
+
+@dataclass
+class LongCatImageEditSamplingParams(SamplingParams):
+    """Defaults for LongCat-Image-Edit (mirrors diffusers LongCatImageEditPipeline).
+
+    Output height/width are derived from the condition image (~1MP), so no
+    defaults are set here. The reference uses an empty negative prompt and
+    plain CFG (no renorm, no prompt rewrite).
+    """
+
+    num_frames: int = 1
+    num_inference_steps: int = 50
+    guidance_scale: float = 4.5
+    negative_prompt: str = ""
+    enable_cfg_renorm: bool = False
+    enable_prompt_rewrite: bool = False
+
+
+@dataclass
+class LongCatImageEditTurboSamplingParams(LongCatImageEditSamplingParams):
+    """LongCat-Image-Edit-Turbo: distilled, 8 steps, CFG disabled."""
+
+    num_inference_steps: int = 8
+    guidance_scale: float = 1.0

python/sglang/multimodal_gen/registry.py

@@ -74,6 +74,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.lingbot_video_moe import (
     LingBotVideoMoEPipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+    LongCatImageEditPipelineConfig,
     LongCatImagePipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig
@@ -150,6 +151,8 @@ from sglang.multimodal_gen.configs.sample.lingbot_world import (
     LingBotWorldSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.longcat_image import (
+    LongCatImageEditSamplingParams,
+    LongCatImageEditTurboSamplingParams,
     LongCatImageSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams
@@ -1327,6 +1330,34 @@ def _register_configs():
         ],
     )
 
+    # LongCat-Image-Edit-Turbo (registered before Edit so its detector wins)
+    register_configs(
+        sampling_param_cls=LongCatImageEditTurboSamplingParams,
+        pipeline_config_cls=LongCatImageEditPipelineConfig,
+        hf_model_paths=[
+            "meituan-longcat/LongCat-Image-Edit-Turbo",
+        ],
+        model_detectors=[
+            lambda hf_id: "longcat" in hf_id.lower()
+            and "edit" in hf_id.lower()
+            and "turbo" in hf_id.lower(),
+        ],
+    )
+
+    # LongCat-Image-Edit
+    register_configs(
+        sampling_param_cls=LongCatImageEditSamplingParams,
+        pipeline_config_cls=LongCatImageEditPipelineConfig,
+        hf_model_paths=[
+            "meituan-longcat/LongCat-Image-Edit",
+        ],
+        model_detectors=[
+            lambda hf_id: "longcat" in hf_id.lower()
+            and "edit" in hf_id.lower()
+            and "turbo" not in hf_id.lower(),
+        ],
+    )
+
 
 _register_configs()
 

python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py

@@ -12,6 +12,10 @@ from transformers import PretrainedConfig
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from sglang.multimodal_gen.configs.models import EncoderConfig
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+    LongCatImageEditPipelineConfig,
+    LongCatImagePipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageEditPipelineConfig,
 )
@@ -869,14 +873,16 @@ class TextEncoderLoader(ComponentLoader):
             with model_device, skip_init_modules():
                 architectures = getattr(model_config, "architectures", [])
                 model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
-                enable_image_understanding = (
-                    True
-                    if isinstance(
-                        server_args.pipeline_config, QwenImageEditPipelineConfig
-                    )
-                    else False
+                enable_image_understanding = isinstance(
+                    server_args.pipeline_config,
+                    (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
                 )
                 model_config.enable_image_understanding = enable_image_understanding
+                # LongCat feeds its padded body to the DiT, so it must mask
+                # padding on the cache-free path; scoped so others are unchanged.
+                model_config.honor_cache_free_padding_mask = isinstance(
+                    server_args.pipeline_config, LongCatImagePipelineConfig
+                )
                 model = model_cls(model_config)
 
             if not isinstance(model, EncoderTensorParallelMixin):

python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py

@@ -443,6 +443,7 @@ class _LongCatSingleAttention(nn.Module):
         self,
         hidden_states: torch.Tensor,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        num_replicated_prefix: int = 0,
         cos_sin_cache: Optional[torch.Tensor] = None,
         positions: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
@@ -464,7 +465,7 @@ class _LongCatSingleAttention(nn.Module):
             positions,
         )
 
-        x = self.attn(q, k, v)
+        x = self.attn(q, k, v, num_replicated_prefix=num_replicated_prefix)
         return x.flatten(2, 3).to(q.dtype)
 
 
@@ -564,6 +565,8 @@ class _SingleTransformerBlock(nn.Module):
         attn_output = self.attn(
             hidden_states=norm_hidden_states,
             image_rotary_emb=image_rotary_emb,
+            # Text is replicated per SP rank; keep it out of the all-to-all.
+            num_replicated_prefix=text_seq_len,
             cos_sin_cache=cos_sin_cache,
             positions=positions,
         )

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -194,6 +194,9 @@ class Qwen2_5_VLAttention(nn.Module):
     def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None):
         super().__init__()
         self.config = config
+        self.honor_cache_free_padding_mask = getattr(
+            config, "honor_cache_free_padding_mask", False
+        )
         self.layer_idx = layer_idx
         if layer_idx is None:
             logger.warning(
@@ -315,14 +318,14 @@ class Qwen2_5_VLAttention(nn.Module):
         query_states = query_states.transpose(1, 2)
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
-        # Diffusion text encoding is cache-free and historically uses the native
-        # causal kernel; its trailing padding is removed during postprocessing.
-        # Cached generation still needs the explicit mask for padded batches.
+        # LongCat masks padding on the cache-free path too; others keep the
+        # original mask-free fast path unchanged.
+        honor_mask = use_cache or self.honor_cache_free_padding_mask
         attn_output = self.attn(
             query_states,
             key_states,
             value_states,
-            attn_mask=attention_mask if use_cache else None,
+            attn_mask=attention_mask if honor_mask else None,
         )
 
         attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
@@ -425,10 +428,27 @@ class Qwen2_5_VLDecoderLayer(nn.Module):
         return hidden_states
 
 
+def _build_causal_padding_mask(
+    attention_mask: torch.Tensor, inputs_embeds: torch.Tensor
+) -> torch.Tensor:
+    """Build a causal+padding bool mask ``[batch, 1, q_len, kv_len]``; True attends."""
+    q_len = inputs_embeds.shape[1]
+    kv_len = attention_mask.shape[-1]
+    device = inputs_embeds.device
+    q_idx = torch.arange(kv_len - q_len, kv_len, device=device).unsqueeze(-1)
+    kv_idx = torch.arange(kv_len, device=device).unsqueeze(0)
+    causal = (kv_idx <= q_idx).unsqueeze(0)
+    padding = attention_mask.to(device=device, dtype=torch.bool).unsqueeze(1)
+    return (causal & padding).unsqueeze(1)
+
+
 class Qwen2_5_VLTextModel(nn.Module):
     def __init__(self, config: PretrainedConfig):
         super().__init__()
         self.config = config
+        self.honor_cache_free_padding_mask = getattr(
+            config, "honor_cache_free_padding_mask", False
+        )
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
@@ -552,6 +572,19 @@ class Qwen2_5_VLTextModel(nn.Module):
                     create_sliding_window_causal_mask(**mask_kwargs)
                 )
 
+            # create_causal_mask returns None for this _attn_implementation, so
+            # build the causal+padding mask here. LongCat-only.
+            if (
+                self.honor_cache_free_padding_mask
+                and causal_mask_mapping["full_attention"] is None
+                and isinstance(attention_mask, torch.Tensor)
+                and attention_mask.dim() == 2
+                and not bool(attention_mask.all())
+            ):
+                causal_mask_mapping["full_attention"] = _build_causal_padding_mask(
+                    attention_mask, inputs_embeds
+                )
+
         hidden_states = inputs_embeds
 
         # decoder layers
@@ -1123,7 +1156,12 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder):
         super().__init__(config)
         enable_image_understanding = config.enable_image_understanding
         generation_config = config.generation_config
+        # LongCat-only; propagate to the text config (see TextEncoderLoader).
+        honor_cache_free_padding_mask = getattr(
+            config, "honor_cache_free_padding_mask", False
+        )
         config = config.arch_config
+        config.text_config.honor_cache_free_padding_mask = honor_cache_free_padding_mask
         self.model = Qwen2_5_VLModel(
             config, enable_image_understanding=enable_image_understanding
         )

python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py

@@ -1,12 +1,19 @@
-"""LongCat-Image pipeline for SGLang."""
+"""LongCat-Image pipelines (T2I and Edit) for SGLang."""
 
 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.pipelines_core.stages import (
+    ImageVAEEncodingStage,
+    InputValidationStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.longcat_image import (
     LongCatPromptRewriteStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.longcat_image_edit import (
+    LongCatImageEditTextEncodingStage,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
@@ -68,4 +75,62 @@ class LongCatImagePipeline(LoRAPipeline, ComposedPipelineBase):
         self.add_standard_decoding_stage()
 
 
-EntryClass = [LongCatImagePipeline]
+class LongCatImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
+    """Pipeline for LongCat-Image-Edit image editing (I2I).
+
+    Mirrors diffusers LongCatImageEditPipeline. The output resolution is
+    derived from the condition image (~1MP, /16); the reference image is
+    VAE-encoded (argmax) and concatenated after the noisy latents; the
+    edit instruction is encoded jointly with the image via Qwen2.5-VL.
+    """
+
+    pipeline_name = "LongCatImageEditPipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "text_processor",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        # 1. Load the condition image, resize it to the calculated output
+        #    resolution, and set batch.height/width (pipeline config hooks).
+        self.add_stage(InputValidationStage())
+
+        # 2. Joint text+image (VL) prompt encoding. Also encodes the negative
+        #    prompt against the same image when CFG is enabled.
+        self.add_stage(
+            LongCatImageEditTextEncodingStage(
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                text_processor=self.get_module("text_processor"),
+                text_encoder_dtype=PRECISION_TO_TYPE[
+                    server_args.pipeline_config.text_encoder_precisions[0]
+                ],
+            )
+        )
+
+        # 3. Reference-image VAE encoding -> packed batch.image_latent, which
+        #    DenoisingStage concatenates after the noisy latents (dim=1).
+        self.add_stage(ImageVAEEncodingStage(vae=self.get_module("vae")))
+
+        # 4. Latent preparation (noise drawn in prompt-embeds dtype).
+        self.add_standard_latent_preparation_stage()
+
+        # 5. Timestep preparation (mu from the packed noisy token count only).
+        self.add_standard_timestep_preparation_stage(
+            prepare_extra_kwargs=[_prepare_mu],
+        )
+
+        # 6. Standard denoising loop; slice_noise_pred drops the reference
+        #    tokens from each prediction before CFG/scheduler step.
+        self.add_standard_denoising_stage()
+
+        # 7. Standard VAE decoding
+        self.add_standard_decoding_stage()
+
+
+EntryClass = [LongCatImagePipeline, LongCatImageEditPipeline]

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -257,6 +257,7 @@ class ComposedPipelineBase(ABC):
                 "QwenImageEditPipeline": {"vae"},
                 "QwenImageEditPlusPipeline": {"vae"},
                 "QwenImageLayeredPipeline": {"vae", "transformer"},
+                "LongCatImageEditPipeline": {"vae"},
                 "GlmImagePipeline": {"vae", "transformer"},
                 "WanImageToVideoPipeline": {"vae"},
                 "WanImageToVideoDmdPipeline": {"vae"},
@@ -1063,6 +1064,10 @@ class ComposedPipelineBase(ABC):
                 main_process_only=True,
             )
 
+        self.component_residency_manager = get_global_component_residency_manager(
+            self, server_args
+        )
+        self.executor.component_residency_manager = self.component_residency_manager
         return self.executor.execute_group_with_profiling(
             self.stages, batches, server_args
         )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py

@@ -0,0 +1,256 @@
+"""Text+image (VL) prompt encoding stage for LongCat-Image-Edit (I2I).
+
+Mirrors diffusers ``LongCatImageEditPipeline._encode_prompt``: the condition
+image (already resized to the output resolution by ``InputValidationStage``)
+is downscaled by 2x and fed to the Qwen2.5-VL vision tower; the prompt body is
+quote-aware tokenized to a fixed 512 tokens; the edit system prefix has its
+``<|image_pad|>`` placeholder expanded to the per-image token count. The
+resulting hidden states are sliced from the ``<|vision_start|>`` token through
+the 512-token body, so the DiT conditioning includes the VL image tokens.
+"""
+
+import PIL.Image
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+    _tokenize_prompt_for_encode,
+)
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+IMAGE_TOKEN = "<|image_pad|>"
+
+# Copied from diffusers LongCatImageEditPipeline.prompt_template_encode_prefix/suffix.
+PROMPT_TEMPLATE_ENCODE_PREFIX = (
+    "<|im_start|>system\nAs an image editing expert, first analyze the content and "
+    "attributes of the input image(s). Then, based on the user's editing instructions, "
+    "clearly and precisely determine how to modify the given image(s), ensuring that "
+    "only the specified parts are altered and all other aspects remain consistent with "
+    "the original(s).<|im_end|>\n<|im_start|>user\n"
+    "<|vision_start|><|image_pad|><|vision_end|>"
+)
+PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+
+
+class LongCatImageEditTextEncodingStage(PipelineStage):
+    """Encode the edit instruction together with the reference image.
+
+    The stage fills ``batch.prompt_embeds`` (and negative counterparts under
+    CFG) with hidden states sliced as ``[vision_start ... vision_end, 512-token
+    body]``. Both branches share the same VL image inputs, matching the
+    diffusers reference which encodes positive and negative prompts against
+    the same ``prompt_image``.
+    """
+
+    deduplicated_output_fields = (
+        "prompt_embeds",
+        "negative_prompt_embeds",
+        "prompt_embeds_mask",
+        "negative_prompt_embeds_mask",
+        "prompt_seq_lens",
+        "negative_prompt_seq_lens",
+    )
+
+    def __init__(
+        self,
+        text_encoder,
+        tokenizer,
+        text_processor,
+        text_encoder_dtype: torch.dtype,
+    ):
+        super().__init__()
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.text_processor = text_processor
+        self.text_encoder_dtype = text_encoder_dtype
+        self._suffix_ids: list[int] | None = None
+
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        return [
+            ComponentUse(
+                stage_name,
+                "text_encoder",
+                target_dtype=self.text_encoder_dtype,
+                memory_intensive=True,
+            ),
+        ]
+
+    def _get_suffix_ids(self) -> list[int]:
+        if self._suffix_ids is None:
+            self._suffix_ids = self.tokenizer(
+                PROMPT_TEMPLATE_ENCODE_SUFFIX, add_special_tokens=False
+            )["input_ids"]
+        return self._suffix_ids
+
+    def _build_expanded_prefix_ids(
+        self, image_grid_thw: torch.Tensor
+    ) -> tuple[list[int], int]:
+        """Expand ``<|image_pad|>`` to the per-image token count and tokenize.
+
+        Returns (prefix_token_ids, prefix_len) where prefix_len is the index of
+        ``<|vision_start|>`` — the slice start for the DiT conditioning.
+        """
+        merge_length = self.text_processor.image_processor.merge_size**2
+        num_image_tokens = int(image_grid_thw.prod().item()) // merge_length
+        text = PROMPT_TEMPLATE_ENCODE_PREFIX
+        while IMAGE_TOKEN in text:
+            text = text.replace(IMAGE_TOKEN, "<|placeholder|>" * num_image_tokens, 1)
+        text = text.replace("<|placeholder|>", IMAGE_TOKEN)
+
+        prefix_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
+        vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+        prefix_len = prefix_ids.index(vision_start_id)
+        return prefix_ids, prefix_len
+
+    def _encode(
+        self,
+        prompt: list[str],
+        pixel_values: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        prefix_ids: list[int],
+        prefix_len: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        body = _tokenize_prompt_for_encode(prompt, self.tokenizer)
+        suffix_ids = self._get_suffix_ids()
+        suffix_len = len(suffix_ids)
+        batch_size = body.input_ids.size(0)
+
+        prefix_ids_t = (
+            torch.tensor(prefix_ids, dtype=body.input_ids.dtype)
+            .unsqueeze(0)
+            .expand(batch_size, -1)
+        )
+        suffix_ids_t = (
+            torch.tensor(suffix_ids, dtype=body.input_ids.dtype)
+            .unsqueeze(0)
+            .expand(batch_size, -1)
+        )
+        prefix_mask_t = torch.ones(
+            batch_size, len(prefix_ids), dtype=body.attention_mask.dtype
+        )
+        suffix_mask_t = torch.ones(
+            batch_size, suffix_len, dtype=body.attention_mask.dtype
+        )
+
+        input_ids = torch.cat((prefix_ids_t, body.input_ids, suffix_ids_t), dim=-1).to(
+            device
+        )
+        attention_mask = torch.cat(
+            (prefix_mask_t, body.attention_mask, suffix_mask_t), dim=-1
+        ).to(device)
+
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            outputs = self.text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
+                output_hidden_states=True,
+                use_cache=False,
+            )
+
+        hidden_states = outputs.hidden_states[-1]
+        # Keep [vision_start ... vision_end, 512-token body].
+        return hidden_states[:, prefix_len:-suffix_len, :]
+
+    @staticmethod
+    def _all_ones_conditioning(prompt_embeds: torch.Tensor):
+        batch_size, seq_len = prompt_embeds.shape[:2]
+        mask = torch.ones(
+            batch_size, seq_len, dtype=torch.bool, device=prompt_embeds.device
+        )
+        return mask, [seq_len] * batch_size
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        device = get_local_torch_device()
+
+        image = batch.condition_image
+        if isinstance(image, list):
+            if len(image) != 1:
+                raise ValueError(
+                    "LongCat-Image-Edit supports exactly one condition image, "
+                    f"got {len(image)}."
+                )
+            image = image[0]
+        if image is None:
+            # Mirrors ImageEncodingStage: nothing to encode without an image
+            # (e.g. warmup requests). Downstream stages will validate inputs.
+            logger.debug(
+                "LongCatImageEditTextEncodingStage skipped: no condition image."
+            )
+            return batch
+
+        prompt = batch.prompt if isinstance(batch.prompt, list) else [batch.prompt]
+        if len(prompt) != 1:
+            raise ValueError(
+                "LongCat-Image-Edit supports a single prompt per request, got "
+                f"{len(prompt)}."
+            )
+
+        # The VL tower sees the condition image at half the output resolution
+        # (matching diffusers: image_processor.resize(image, h // 2, w // 2)).
+        prompt_image = image.resize(
+            (int(batch.width) // 2, int(batch.height) // 2),
+            PIL.Image.Resampling.LANCZOS,
+        )
+        vl_inputs = self.text_processor.image_processor(
+            images=prompt_image, return_tensors="pt"
+        )
+        image_grid_thw = vl_inputs["image_grid_thw"].to(device)
+        prefix_ids, prefix_len = self._build_expanded_prefix_ids(image_grid_thw)
+
+        with self.use_declared_component(
+            component_name="text_encoder",
+            module=self.text_encoder,
+        ) as text_encoder:
+            assert text_encoder is not None
+            self.text_encoder = text_encoder
+
+            pixel_values = vl_inputs["pixel_values"].to(
+                device=device, dtype=self.text_encoder_dtype
+            )
+
+            prompt_embeds = self._encode(
+                prompt, pixel_values, image_grid_thw, prefix_ids, prefix_len, device
+            )
+            batch.prompt_embeds.append(prompt_embeds)
+            mask, seq_lens = self._all_ones_conditioning(prompt_embeds)
+            batch.prompt_embeds_mask = [mask]
+            batch.prompt_seq_lens = [seq_lens]
+
+            if batch.do_classifier_free_guidance:
+                negative_prompt = batch.negative_prompt or ""
+                if isinstance(negative_prompt, list):
+                    negative_prompt = negative_prompt[:1]
+                else:
+                    negative_prompt = [negative_prompt]
+                negative_prompt_embeds = self._encode(
+                    negative_prompt,
+                    pixel_values,
+                    image_grid_thw,
+                    prefix_ids,
+                    prefix_len,
+                    device,
+                )
+                batch.negative_prompt_embeds.append(negative_prompt_embeds)
+                neg_mask, neg_seq_lens = self._all_ones_conditioning(
+                    negative_prompt_embeds
+                )
+                batch.negative_prompt_embeds_mask = [neg_mask]
+                batch.negative_prompt_seq_lens = [neg_seq_lens]
+
+        return batch

python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py

@@ -0,0 +1,157 @@
+"""Unit tests for LongCat-Image-Edit pipeline config hooks (CPU only)."""
+
+import types
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
+from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import (
+    LongCatImageEditPipelineConfig,
+    LongCatImagePipelineConfig,
+    _calculate_edit_dimensions,
+)
+
+
+@pytest.fixture
+def edit_config():
+    return LongCatImageEditPipelineConfig()
+
+
+def _make_batch(**kwargs):
+    return types.SimpleNamespace(**kwargs)
+
+
+def test_edit_dimensions_match_diffusers_formula():
+    import math
+
+    for w, h in [(1000, 1500), (1024, 1024), (1920, 1080), (333, 777)]:
+        ratio = w / h
+        got_w, got_h = _calculate_edit_dimensions(1024 * 1024, ratio)
+        # diffusers reference (ceil to /16)
+        ref_w = math.sqrt(1024 * 1024 * ratio)
+        ref_h = ref_w / ratio
+        ref_w = ref_w if ref_w % 16 == 0 else (ref_w // 16 + 1) * 16
+        ref_h = ref_h if ref_h % 16 == 0 else (ref_h // 16 + 1) * 16
+        assert (got_w, got_h) == (int(ref_w), int(ref_h))
+        assert got_w % 16 == 0 and got_h % 16 == 0
+
+
+def test_edit_config_task_type_and_generator(edit_config):
+    assert edit_config.task_type == ModelTaskType.I2I
+    assert edit_config.generator_device == "cpu"
+
+
+def test_slice_noise_pred_drops_reference_tokens(edit_config):
+    latents = torch.zeros(1, 100, 64)
+    noise = torch.arange(200, dtype=torch.float32).view(1, 200, 1).expand(1, 200, 64)
+    sliced = edit_config.slice_noise_pred(noise, latents)
+    assert sliced.shape == (1, 100, 64)
+    assert torch.equal(sliced, noise[:, :100])
+
+
+def test_maybe_prepare_latent_ids_returns_none(edit_config):
+    assert edit_config.maybe_prepare_latent_ids(torch.zeros(1, 16, 8, 8)) is None
+
+
+def test_preprocess_vae_encode_squeezes_frames_dim(edit_config):
+    image_5d = torch.zeros(1, 3, 1, 64, 64)
+    assert edit_config.preprocess_vae_encode(image_5d, None).shape == (1, 3, 64, 64)
+    image_4d = torch.zeros(1, 3, 64, 64)
+    assert edit_config.preprocess_vae_encode(image_4d, None).shape == (1, 3, 64, 64)
+
+
+def test_postprocess_image_latent_packs_2x2(edit_config):
+    latent = torch.randn(1, 16, 8, 8)
+    batch = _make_batch(batch_size=1)
+    packed = edit_config.postprocess_image_latent(latent, batch)
+    assert packed.shape == (1, 16, 64)  # (8//2 * 8//2, 16*4)
+
+
+def test_edit_img_ids_modalities_and_offset(edit_config):
+    # 1024x1024 -> latent 128x128 -> packed grid 64x64 -> 4096 tokens per image
+    batch = _make_batch(height=1024, width=1024)
+    num_token = 600  # text length including VL image tokens
+    img_ids = edit_config._edit_img_ids(batch, num_token, torch.device("cpu"))
+    assert img_ids.shape == (2 * 64 * 64, 3)
+    noisy_ids, ref_ids = img_ids[: 64 * 64], img_ids[64 * 64 :]
+    assert torch.all(noisy_ids[:, 0] == 1)
+    assert torch.all(ref_ids[:, 0] == 2)
+    # same grid layout, offset by text length
+    assert torch.equal(noisy_ids[:, 1:], ref_ids[:, 1:])
+    assert noisy_ids[0, 1].item() == num_token
+    assert noisy_ids[0, 2].item() == num_token
+    assert noisy_ids[-1, 1].item() == num_token + 63
+    assert noisy_ids[-1, 2].item() == num_token + 63
+
+
+def test_prepare_pos_cond_kwargs_uses_prompt_embeds_length(edit_config):
+    prompt_embeds = torch.zeros(1, 700, 3072)
+    batch = _make_batch(height=512, width=512, prompt_embeds=[prompt_embeds])
+    kwargs = edit_config.prepare_pos_cond_kwargs(
+        batch, torch.device("cpu"), None, torch.float32
+    )
+    txt_ids, img_ids = kwargs["txt_ids"], kwargs["img_ids"]
+    assert txt_ids.shape == (700, 3)
+    assert torch.all(txt_ids[:, 0] == 0)
+    # 512 -> latent 64 -> packed grid 32x32
+    assert img_ids.shape == (2 * 32 * 32, 3)
+    assert img_ids[0, 1].item() == 700
+
+
+def test_t2i_config_unchanged():
+    t2i = LongCatImagePipelineConfig()
+    assert t2i.task_type == ModelTaskType.T2I
+    # T2I still slices nothing and builds latent ids at preparation time
+    noise = torch.zeros(1, 10, 64)
+    assert t2i.slice_noise_pred(noise, torch.zeros(1, 5, 64)).shape == (1, 10, 64)
+    ids = t2i.maybe_prepare_latent_ids(torch.zeros(1, 16, 8, 8))
+    assert ids is not None and ids.shape == (16, 3)
+
+
+def test_expand_conditioning_repeats_embeds_for_num_outputs():
+    # Single prompt, num_outputs_per_prompt=2: text encoding produces per-prompt
+    # (batch 1) embeds while latents are built at batch 2, so the conditioning
+    # must be repeated to match or the DiT sees mismatched batch dims.
+    config = LongCatImageEditPipelineConfig()
+    batch = _make_batch(
+        prompt=["edit it"],
+        num_outputs_per_prompt=2,
+        prompt_embeds=[torch.randn(1, 850, 3584)],
+        negative_prompt_embeds=[torch.randn(1, 850, 3584)],
+        prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+        negative_prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+        prompt_seq_lens=[[850]],
+        negative_prompt_seq_lens=[[850]],
+    )
+    pos0 = batch.prompt_embeds[0]
+
+    config.expand_conditioning_to_sample_batch(batch)
+
+    assert batch.prompt_embeds[0].shape == (2, 850, 3584)
+    assert batch.negative_prompt_embeds[0].shape == (2, 850, 3584)
+    assert batch.prompt_embeds_mask[0].shape == (2, 850)
+    assert batch.negative_prompt_embeds_mask[0].shape == (2, 850)
+    assert batch.prompt_seq_lens[0] == [850, 850]
+    assert batch.negative_prompt_seq_lens[0] == [850, 850]
+    # Each sample is the original prompt, not garbage.
+    assert torch.equal(batch.prompt_embeds[0][0], pos0[0])
+    assert torch.equal(batch.prompt_embeds[0][1], pos0[0])
+
+
+def test_expand_conditioning_noop_for_single_output():
+    config = LongCatImagePipelineConfig()
+    pe = torch.randn(1, 850, 3584)
+    batch = _make_batch(
+        prompt=["a photo"],
+        num_outputs_per_prompt=1,
+        prompt_embeds=[pe],
+        negative_prompt_embeds=None,
+        prompt_embeds_mask=[torch.ones(1, 850, dtype=torch.bool)],
+        negative_prompt_embeds_mask=None,
+        prompt_seq_lens=[[850]],
+        negative_prompt_seq_lens=None,
+    )
+    config.expand_conditioning_to_sample_batch(batch)
+    assert batch.prompt_embeds[0] is pe  # untouched
+    assert batch.prompt_embeds[0].shape == (1, 850, 3584)

python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py

@@ -148,7 +148,7 @@ def test_text_mlp_uses_single_rank_when_intermediate_size_is_not_tp_divisible(
     assert isinstance(layer.mlp.down_proj, ReplicatedLinear)
 
 
-def test_explicit_attention_mask_is_limited_to_cached_generation(monkeypatch):
+def test_explicit_attention_mask_is_honored_without_a_cache(monkeypatch):
     attention = Qwen2_5_VLAttention.__new__(Qwen2_5_VLAttention)
     nn.Module.__init__(attention)
     attention.q_proj = nn.Identity()
@@ -174,9 +174,19 @@ def test_explicit_attention_mask_is_limited_to_cached_generation(monkeypatch):
         "position_ids": torch.zeros(3, 1, 2, dtype=torch.long),
     }
 
+    # LongCat opts into masking the padded body on the cache-free path.
+    attention.honor_cache_free_padding_mask = True
     attention(**kwargs, use_cache=False)
     attention(**kwargs, use_cache=True)
+    assert attention.attn.masks[0] is explicit_mask
+    assert attention.attn.masks[1] is explicit_mask
 
+    # Every other pipeline keeps the original behavior: mask dropped when
+    # cache-free, honored only under cached generation.
+    attention.attn.masks.clear()
+    attention.honor_cache_free_padding_mask = False
+    attention(**kwargs, use_cache=False)
+    attention(**kwargs, use_cache=True)
     assert attention.attn.masks[0] is None
     assert attention.attn.masks[1] is explicit_mask