diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx index 2426e96c4..2a825570e 100644 --- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/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 LongCat-Image -
meituan-longcat/LongCat-Image
+
meituan-longcat/LongCat-Imagemeituan-longcat/LongCat-Image-Editmeituan-longcat/LongCat-Image-Edit-Turbo
SD3 / SD3.5 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py index 8f620e9c0..2d0daf4a9 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py +++ b/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) diff --git a/python/sglang/multimodal_gen/configs/sample/longcat_image.py b/python/sglang/multimodal_gen/configs/sample/longcat_image.py index c9999436b..91d35b5c2 100644 --- a/python/sglang/multimodal_gen/configs/sample/longcat_image.py +++ b/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 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index e78679527..9d09359ad 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index f86f80204..c7383ccf2 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py index 9d22fe7fa..f7c2dbc08 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index d6bf3dbc4..ab8b74f0e 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/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 ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py b/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py index 87222a1e7..b843026ad 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py +++ b/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] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 5c6b0c201..3f0ada45f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/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 ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image_edit.py new file mode 100644 index 000000000..5a3532339 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py b/python/sglang/multimodal_gen/test/unit/test_longcat_image_edit_config.py new file mode 100644 index 000000000..3069a8630 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py b/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py index b5d9d20ee..13b6f444a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen2_5vl_generation.py +++ b/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