diff --git a/python/sglang/multimodal_gen/configs/models/dits/longcat_image.py b/python/sglang/multimodal_gen/configs/models/dits/longcat_image.py new file mode 100644 index 000000000..696327dbd --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/longcat_image.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class LongCatImageArchConfig(DiTArchConfig): + patch_size: int = 1 + in_channels: int = 64 # packed: 16 * 4 + num_layers: int = 19 + num_single_layers: int = 38 + attention_head_dim: int = 128 + num_attention_heads: int = 24 + joint_attention_dim: int = 3584 + pooled_projection_dim: int = 3584 + axes_dims_rope: list = field(default_factory=lambda: [16, 56, 56]) + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = 16 # unpacked channels + + +@dataclass +class LongCatImageDitConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=LongCatImageArchConfig) + + prefix: str = "longcat_image" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/longcat_image.py b/python/sglang/multimodal_gen/configs/models/vaes/longcat_image.py new file mode 100644 index 000000000..87cf594bf --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/longcat_image.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class LongCatImageVAEArchConfig(VAEArchConfig): + spatial_compression_ratio: int = 8 + vae_scale_factor: int = 8 + # scaling_factor and shift_factor come from the model's config at runtime + + +@dataclass +class LongCatImageVAEConfig(VAEConfig): + arch_config: LongCatImageVAEArchConfig = field( + default_factory=LongCatImageVAEArchConfig + ) + + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def get_vae_scale_factor(self): + return self.arch_config.vae_scale_factor diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py new file mode 100644 index 000000000..f9cad23ab --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py @@ -0,0 +1,422 @@ +from dataclasses import dataclass, field +from typing import Callable + +import numpy as np +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.longcat_image import ( + LongCatImageDitConfig, +) +from sglang.multimodal_gen.configs.models.vaes.longcat_image import ( + LongCatImageVAEConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, + TextConditioningOutput, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Encode-side prompt length (governs tokenization truncation/padding and the +# image latent position-id start offset). Mirrors diffusers' +# LongCatImagePipeline.tokenizer_max_length. +TOKENIZER_MAX_LENGTH = 512 + +# Fixed chat-template wrappers prepended/appended around the encoded prompt. +# Mirrors diffusers LongCatImagePipeline._encode_prompt. +ENCODE_PREFIX_STR = ( + "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt " + "based on an image content, suitable for input to a text-to-image model.<|im_end|>\n" + "<|im_start|>user\n" +) +ENCODE_SUFFIX_STR = "<|im_end|>\n<|im_start|>assistant\n" + + +def _split_quotation(prompt, quote_pairs=None): + """Split prompt on quoted substrings, returning list of (text, is_quoted) tuples.""" + import re + + word_internal_quote_pattern = re.compile(r"[a-zA-Z]+'[a-zA-Z]+") + matches = word_internal_quote_pattern.findall(prompt) + mapping = [] + + for i, word_src in enumerate(set(matches)): + word_tgt = "longcat_$##$_longcat" * (i + 1) + prompt = prompt.replace(word_src, word_tgt) + mapping.append([word_src, word_tgt]) + + if quote_pairs is None: + quote_pairs = [ + ("'", "'"), + ('"', '"'), + ("‘", "’"), + ("“", "”"), + ] + pattern = "|".join( + [ + re.escape(q1) + r"[^" + re.escape(q1 + q2) + r"]*?" + re.escape(q2) + for q1, q2 in quote_pairs + ] + ) + parts = re.split(f"({pattern})", prompt) + + result = [] + for part in parts: + for word_src, word_tgt in mapping: + part = part.replace(word_tgt, word_src) + if re.match(pattern, part): + if len(part): + result.append((part, True)) + else: + if len(part): + result.append((part, False)) + return result + + +def _prepare_pos_ids( + modality_id=0, + token_type="text", + start=(0, 0), + num_token=None, + height=None, + width=None, +): + if token_type == "text": + assert num_token + pos_ids = torch.zeros(num_token, 3) + pos_ids[..., 0] = modality_id + pos_ids[..., 1] = torch.arange(num_token) + start[0] + pos_ids[..., 2] = torch.arange(num_token) + start[1] + elif token_type == "image": + assert height and width + pos_ids = torch.zeros(height, width, 3) + pos_ids[..., 0] = modality_id + pos_ids[..., 1] = pos_ids[..., 1] + torch.arange(height)[:, None] + start[0] + pos_ids[..., 2] = pos_ids[..., 2] + torch.arange(width)[None, :] + start[1] + pos_ids = pos_ids.reshape(height * width, 3) + else: + raise KeyError( + f'Unknown token_type {token_type}, only support "text" or "image".' + ) + return pos_ids + + +def _tokenize_prompt_for_encode(prompt, tokenizer): + """Quote-aware tokenization mirroring diffusers LongCatImagePipeline._encode_prompt. + + Quoted substrings are tokenized character-by-character; unquoted substrings + are tokenized whole. Truncated/padded to TOKENIZER_MAX_LENGTH. Returns the + padded (input_ids, attention_mask) for the prompt body (without prefix/suffix). + """ + if isinstance(prompt, str): + prompt = [prompt] + + batch_all_tokens = [] + for each_prompt in prompt: + all_tokens = [] + for clean_prompt_sub, matched in _split_quotation(each_prompt): + if matched: + # Intentional: tokenize each character in quoted text individually, + # mirroring diffusers LongCatImagePipeline._encode_prompt behavior. + for sub_word in clean_prompt_sub: + tokens = tokenizer(sub_word, add_special_tokens=False)["input_ids"] + all_tokens.extend(tokens) + else: + tokens = tokenizer(clean_prompt_sub, add_special_tokens=False)[ + "input_ids" + ] + all_tokens.extend(tokens) + + if len(all_tokens) > TOKENIZER_MAX_LENGTH: + logger.warning( + "Prompt truncated: max_sequence_length=%d, input_token_nums=%d", + TOKENIZER_MAX_LENGTH, + len(all_tokens), + ) + all_tokens = all_tokens[:TOKENIZER_MAX_LENGTH] + batch_all_tokens.append(all_tokens) + + text_tokens_and_mask = tokenizer.pad( + {"input_ids": batch_all_tokens}, + max_length=TOKENIZER_MAX_LENGTH, + padding="max_length", + return_attention_mask=True, + return_tensors="pt", + ) + return text_tokens_and_mask + + +def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # VAE applies 8x compression, plus 2x packing factor + h = 2 * (int(height) // (vae_scale_factor * 2)) + w = 2 * (int(width) // (vae_scale_factor * 2)) + + latents = latents.view(batch_size, h // 2, w // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + latents = latents.reshape(batch_size, channels // (2 * 2), h, w) + return latents + + +def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view( + batch_size, num_channels_latents, height // 2, 2, width // 2, 2 + ) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape( + batch_size, (height // 2) * (width // 2), num_channels_latents * 4 + ) + return latents + + +def _calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + return image_seq_len * m + b + + +def longcat_postprocess_text(outputs, text_inputs, pipeline_config): + """Slice hidden states to drop the encode prefix/suffix tokens. + + The tokenizer (via `LongCatImagePipelineConfig.tokenize_prompt`) wraps each + prompt with a fixed system prefix and an assistant suffix. The DiT must + receive only the 512-token prompt body, so we slice + `hidden_states[-1][:, prefix_len:-suffix_len, :]`. A mask aligned to the + sliced length is returned so `build_text_conditioning_mask`'s length check + (raw_mask len vs embed seq len) does not raise on the unsliced + prefix+suffix mask. + + LongCat feeds the full 512-token body (including padding) to the DiT, so the + mask is all-ones and seq_lens are all 512 — preserving the pre-refactor + behavior where the DiT attended over the entire padded prompt body. + """ + prefix_len = pipeline_config._encode_prefix_len + suffix_len = pipeline_config._encode_suffix_len + + hidden_states = outputs.hidden_states[-1] + prompt_embeds = hidden_states[:, prefix_len:-suffix_len, :] + + seq_len = prompt_embeds.shape[1] + batch_size = prompt_embeds.shape[0] + prompt_embeds_mask = torch.ones( + batch_size, seq_len, dtype=torch.bool, device=prompt_embeds.device + ) + prompt_seq_lens = [seq_len] * batch_size + return TextConditioningOutput( + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + prompt_seq_lens=prompt_seq_lens, + ) + + +@dataclass +class LongCatImageEncoderConfig(EncoderConfig): + """Encoder config for the in-stage-loaded HF Qwen2.5-VL text encoder. + + The encoder weights are loaded by `LongCatPromptRewriteStage` (not via + `TextEncoderLoader`), so this config only supplies the fields the standard + `TextEncodingStage` reads — primarily `tokenizer_kwargs`. + """ + + tokenizer_kwargs: dict = field(default_factory=lambda: {}) + + +@dataclass +class LongCatImagePipelineConfig(ImagePipelineConfig): + """Configuration for the LongCat-Image T2I pipeline.""" + + task_type: ModelTaskType = ModelTaskType.T2I + + vae_precision: str = "bf16" + should_use_guidance: bool = True + vae_tiling: bool = False + vae_sp: bool = False + enable_autocast: bool = False + + dit_config: DiTConfig = field(default_factory=LongCatImageDitConfig) + vae_config: VAEConfig = field(default_factory=LongCatImageVAEConfig) + + # The Qwen2.5-VL text encoder (~7B) is loaded in bf16; the encoder is loaded + # in-stage by LongCatPromptRewriteStage, not via TextEncoderLoader. + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (LongCatImageEncoderConfig(),) + ) + postprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (longcat_postprocess_text,) + ) + + # --- LatentPreparationStage hooks --- + + def prepare_latent_shape(self, batch, batch_size, num_frames): + vae_scale_factor = self.vae_config.get_vae_scale_factor() + # LongCat packs 2x2 patches: effective spatial resolution after packing + h = 2 * (int(batch.height) // (vae_scale_factor * 2)) + w = 2 * (int(batch.width) // (vae_scale_factor * 2)) + num_channels_latents = self.dit_config.arch_config.num_channels_latents + # Unpacked shape — maybe_pack_latents will fold into tokens + return (batch_size, num_channels_latents, h, w) + + def maybe_pack_latents(self, latents, batch_size, batch): + num_channels_latents = self.dit_config.arch_config.num_channels_latents + _, _, h, w = latents.shape + return _pack_latents(latents, batch_size, num_channels_latents, h, w) + + def maybe_prepare_latent_ids(self, latents): + # latents shape after packing: [B, (h//2)*(w//2), C*4] + # We need h//2 and w//2 — derive from the unpacked shape stored on the config. + # latents is still unpacked here (called before maybe_pack_latents in LatentPreparationStage) + _, _, h, w = latents.shape + return _prepare_pos_ids( + modality_id=1, + token_type="image", + start=(TOKENIZER_MAX_LENGTH, TOKENIZER_MAX_LENGTH), + height=h // 2, + width=w // 2, + ) + + def get_latent_dtype(self, prompt_dtype: torch.dtype) -> torch.dtype: + # Generate in float32 then cast to bfloat16, matching diffusers behavior. + return torch.float32 + + # --- TextEncodingStage hooks --- + + def _ensure_encode_prefix_suffix(self, tokenizer): + """Lazily tokenize the fixed encode prefix/suffix (tokenizer unavailable + at config construction time). Cached on the config instance.""" + if not hasattr(self, "_encode_prefix_ids"): + self._encode_prefix_ids = tokenizer( + ENCODE_PREFIX_STR, add_special_tokens=False + )["input_ids"] + self._encode_suffix_ids = tokenizer( + ENCODE_SUFFIX_STR, add_special_tokens=False + )["input_ids"] + self._encode_prefix_len = len(self._encode_prefix_ids) + self._encode_suffix_len = len(self._encode_suffix_ids) + + def tokenize_prompt(self, prompt, tokenizer, tok_kwargs): + """Quote-aware tokenization + fixed prefix/suffix wrapping. + + Mirrors diffusers LongCatImagePipeline._encode_prompt: quoted substrings + are tokenized character-by-character, the body is truncated/padded to + TOKENIZER_MAX_LENGTH, then a fixed system prefix and assistant suffix are + concatenated onto every sequence (with all-ones masks). Returns a + BatchEncoding-like dict with `input_ids` / `attention_mask` of length + `prefix_len + TOKENIZER_MAX_LENGTH + suffix_len`. + """ + from transformers import BatchEncoding + + self._ensure_encode_prefix_suffix(tokenizer) + + body = _tokenize_prompt_for_encode(prompt, tokenizer) + + prefix_len = self._encode_prefix_len + suffix_len = self._encode_suffix_len + batch_size = body.input_ids.size(0) + + prefix_ids_t = ( + torch.tensor(self._encode_prefix_ids, dtype=body.input_ids.dtype) + .unsqueeze(0) + .expand(batch_size, -1) + ) + suffix_ids_t = ( + torch.tensor(self._encode_suffix_ids, dtype=body.input_ids.dtype) + .unsqueeze(0) + .expand(batch_size, -1) + ) + prefix_mask_t = torch.ones( + batch_size, prefix_len, 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) + attention_mask = torch.cat( + (prefix_mask_t, body.attention_mask, suffix_mask_t), dim=-1 + ) + # Return a BatchEncoding so TextEncodingStage can call `.to(device)` on + # the result and index input_ids / attention_mask like a tokenizer output. + return BatchEncoding( + data={"input_ids": input_ids, "attention_mask": attention_mask} + ) + + # --- TimestepPreparationStage hook --- + + def prepare_sigmas(self, sigmas, num_inference_steps): + if sigmas is None: + sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) + return sigmas + + def get_pos_prompt_embeds(self, batch): + return batch.prompt_embeds[0] + + def get_neg_prompt_embeds(self, batch): + return batch.negative_prompt_embeds[0] + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + # txt_ids depend only on the prompt length (fixed at TOKENIZER_MAX_LENGTH); + # img_ids are the latent position ids set by LatentPreparationStage. + # image_rotary_emb is NOT precomputed — the DiT computes it on the fly from + # txt_ids + img_ids (matching diffusers' transformer, which calls + # self.pos_embed internally). + 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": batch.latent_ids, + } + + 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": batch.latent_ids, + } + + def get_decode_scale_and_shift(self, device, dtype, vae): + # scaling_factor/shift_factor live on the VAE runtime config (not the + # arch config). AutoencoderKL always has them; access directly so a + # missing field raises instead of silently decoding with 1.0/0.0. + return vae.config.scaling_factor, vae.config.shift_factor + + def post_denoising_loop(self, latents, batch): + vae_scale_factor = self.vae_config.get_vae_scale_factor() + latents = _unpack_latents(latents, batch.height, batch.width, vae_scale_factor) + # Add frames dimension for DecodingStage compatibility: [B, C, H, W] -> [B, C, 1, H, W] + latents = latents.unsqueeze(2) + return latents + + def preprocess_decoding(self, latents, server_args=None, vae=None): + """Remove frames dimension before VAE decode: [B, C, 1, H, W] -> [B, C, H, W].""" + if latents.dim() == 5 and latents.shape[2] == 1: + latents = latents.squeeze(2) + return latents + + def postprocess_cfg_noise( + self, + batch, + noise_pred: torch.Tensor, + noise_pred_cond: torch.Tensor, + ) -> torch.Tensor: + enable_cfg_renorm = getattr(batch, "enable_cfg_renorm", True) + cfg_renorm_min = getattr(batch, "cfg_renorm_min", 0.0) + if not enable_cfg_renorm: + return noise_pred + cond_norm = torch.norm(noise_pred_cond, dim=-1, keepdim=True) + 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 diff --git a/python/sglang/multimodal_gen/configs/sample/longcat_image.py b/python/sglang/multimodal_gen/configs/sample/longcat_image.py new file mode 100644 index 000000000..c9999436b --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/longcat_image.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class LongCatImageSamplingParams(SamplingParams): + num_inference_steps: int = 50 + guidance_scale: float = 4.5 + height: int = 1024 + width: int = 1024 + # Override base class defaults to enable LongCat-specific features by default + enable_cfg_renorm: bool = True + enable_prompt_rewrite: bool = True diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 7390f8f62..33ab21614 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -201,6 +201,11 @@ class SamplingParams: progressive_levels: int = 1 progressive_delta: float = 0.01 + # LongCat-Image parameters + enable_cfg_renorm: bool = False + cfg_renorm_min: float = 0.0 + enable_prompt_rewrite: bool = False + # TeaCache parameters enable_teacache: bool = False teacache_params: Any = ( @@ -952,6 +957,23 @@ class SamplingParams: help="Spectrum tau normalization horizon.", ) + # LongCat-Image parameters + add_argument( + "--enable-cfg-renorm", + action=StoreBoolean, + help="Enable CFG renormalization for LongCat-Image (default: false).", + ) + add_argument( + "--cfg-renorm-min", + type=float, + help="Minimum CFG renorm scale for LongCat-Image (default: 0.0).", + ) + add_argument( + "--enable-prompt-rewrite", + action=StoreBoolean, + help="Enable prompt rewriting via Qwen2.5-VL before encoding for LongCat-Image (default: false).", + ) + # profiling add_argument( "--profile", diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index df4931c4d..8e9f360da 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -73,6 +73,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.krea2 import Krea2PipelineCo from sglang.multimodal_gen.configs.pipeline_configs.lingbot_video_moe import ( LingBotVideoMoEPipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import ( + LongCatImagePipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, @@ -145,6 +148,9 @@ from sglang.multimodal_gen.configs.sample.lingbot_video_moe import ( from sglang.multimodal_gen.configs.sample.lingbot_world import ( LingBotWorldSamplingParams, ) +from sglang.multimodal_gen.configs.sample.longcat_image import ( + LongCatImageSamplingParams, +) from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams from sglang.multimodal_gen.configs.sample.ltx_2 import ( LTX2SamplingParams, @@ -1211,6 +1217,18 @@ def _register_configs(): ], ) + # LongCat-Image + register_configs( + sampling_param_cls=LongCatImageSamplingParams, + pipeline_config_cls=LongCatImagePipelineConfig, + hf_model_paths=[ + "meituan-longcat/LongCat-Image", + ], + model_detectors=[ + lambda hf_id: "longcat" in hf_id.lower() and "edit" not in hf_id.lower(), + ], + ) + _register_configs() diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index b1dbf3766..21f9ff7bc 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -501,7 +501,7 @@ class ImageProcessorLoader(ComponentLoader): class AutoProcessorLoader(ComponentLoader): """Loader for auto processor.""" - component_names = ["processor"] + component_names = ["processor", "text_processor"] expected_library = "transformers" def load_customized( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py new file mode 100644 index 000000000..7257ade0b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py @@ -0,0 +1,710 @@ +# Copied and adapted from: https://github.com/huggingface/diffusers +# main/src/diffusers/models/transformers/transformer_longcat_image.py +"""LongCat-Image Transformer2D model for SGLang. + +Implements the LongCatImageTransformer2DModel architecture from diffusers +adapted for SGLang's inference pipeline. + +Key adaptation: the transformer accepts raw scheduler timesteps (in [0, 1000]) +and feeds them directly to the timestep embedder. The diffusers pipeline passes +``timestep / 1000`` to its transformer (which then multiplies by 1000 internally); +SGLang's DenoisingStage passes the raw timestep instead, so the value reaching +the embedder is identical and no division is needed here. + +Attention alignment: uses USPAttention (FA3/FA4 on Hopper/Blackwell) with +SGLang fused RMSNorm (apply_qk_norm). RoPE is applied separately via +diffusers apply_rotary_emb because LongCat's axes_dims_rope=[16,56,56] +sums to head_dim=128 (full rotation), which is incompatible with flashinfer's +cos_sin_cache format that requires rotary_dim <= head_dim. +""" + +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +from diffusers.models.embeddings import ( + TimestepEmbedding, + Timesteps, + apply_rotary_emb, + get_1d_rotary_pos_embed, +) +from diffusers.models.normalization import ( + AdaLayerNormContinuous, + AdaLayerNormZero, + AdaLayerNormZeroSingle, +) + +from sglang.multimodal_gen.runtime.distributed import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm +from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# FFN +# --------------------------------------------------------------------------- + + +class _LongCatFFN(nn.Module): + """TP-parallel FFN matching diffusers FeedForward(activation_fn="gelu-approximate"). + + Weight names mirror the diffusers checkpoint layout so loading works without remapping: + net.0.proj -> ColumnParallelLinear (project in, dim -> inner_dim) + net.2 -> RowParallelLinear (project out, inner_dim -> dim) + """ + + def __init__(self, dim: int, inner_dim: int, bias: bool = True, prefix: str = ""): + super().__init__() + self.net = nn.ModuleList( + [ + # net.0: GELU activation wrapper — only the inner proj is a parameter + nn.ModuleDict( + { + "proj": ColumnParallelLinear( + dim, + inner_dim, + bias=bias, + gather_output=False, + prefix=f"{prefix}.net.0.proj", + ) + } + ), + nn.Dropout( + 0.0 + ), # net.1: dummy dropout, matches diffusers checkpoint layout + RowParallelLinear( + inner_dim, + dim, + bias=bias, + input_is_parallel=True, + prefix=f"{prefix}.net.2", + ), + ] + ) + self.act = nn.GELU(approximate="tanh") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.net[0]["proj"](hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.net[2](hidden_states) + return hidden_states + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class _LongCatJointAttention(nn.Module): + """Double-stream (joint) attention for _TransformerBlock. + + img and txt tokens are projected separately, QK-norm applied via SGLang + fused kernel, RoPE applied via diffusers apply_rotary_emb (supports full + head_dim rotation), then concatenated (txt first) before USPAttention. + + TP: Q/K/V and add_q/k/v use ColumnParallelLinear (heads sharded across TP ranks). + Output projections use RowParallelLinear (all-reduce after matmul). + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + bias: bool = True, + eps: float = 1e-6, + prefix: str = "", + ): + super().__init__() + tp_size = get_tp_world_size() + self.num_local_heads = num_attention_heads // tp_size + assert ( + num_attention_heads % tp_size == 0 + ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + self.head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + self.norm_q = RMSNorm(attention_head_dim, eps=eps) + self.norm_k = RMSNorm(attention_head_dim, eps=eps) + self.norm_added_q = RMSNorm(attention_head_dim, eps=eps) + self.norm_added_k = RMSNorm(attention_head_dim, eps=eps) + + self.to_q = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_q" + ) + self.to_k = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_k" + ) + self.to_v = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_v" + ) + self.add_q_proj = ColumnParallelLinear( + dim, + inner_dim, + bias=bias, + gather_output=False, + prefix=f"{prefix}.add_q_proj", + ) + self.add_k_proj = ColumnParallelLinear( + dim, + inner_dim, + bias=bias, + gather_output=False, + prefix=f"{prefix}.add_k_proj", + ) + self.add_v_proj = ColumnParallelLinear( + dim, + inner_dim, + bias=bias, + gather_output=False, + prefix=f"{prefix}.add_v_proj", + ) + + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + inner_dim, + dim, + bias=bias, + input_is_parallel=True, + prefix=f"{prefix}.to_out.0", + ) + ] + ) + self.to_add_out = RowParallelLinear( + inner_dim, + dim, + bias=bias, + input_is_parallel=True, + prefix=f"{prefix}.to_add_out", + ) + + self.attn = USPAttention( + num_heads=self.num_local_heads, + head_size=attention_head_dim, + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + txt_seq_len = encoder_hidden_states.shape[1] + + q, _ = self.to_q(hidden_states) + k, _ = self.to_k(hidden_states) + v, _ = self.to_v(hidden_states) + q = q.unflatten(-1, (self.num_local_heads, self.head_dim)) + k = k.unflatten(-1, (self.num_local_heads, self.head_dim)) + v = v.unflatten(-1, (self.num_local_heads, self.head_dim)) + + eq, _ = self.add_q_proj(encoder_hidden_states) + ek, _ = self.add_k_proj(encoder_hidden_states) + ev, _ = self.add_v_proj(encoder_hidden_states) + eq = eq.unflatten(-1, (self.num_local_heads, self.head_dim)) + ek = ek.unflatten(-1, (self.num_local_heads, self.head_dim)) + ev = ev.unflatten(-1, (self.num_local_heads, self.head_dim)) + + # SGLang fused QK-norm + q, k = apply_qk_norm(q, k, self.norm_q, self.norm_k, self.head_dim) + eq, ek = apply_qk_norm( + eq, ek, self.norm_added_q, self.norm_added_k, self.head_dim + ) + + # Concatenate: txt first, then img (matches diffusers convention) + q = torch.cat([eq, q], dim=1) + k = torch.cat([ek, k], dim=1) + v = torch.cat([ev, v], dim=1) + + # RoPE applied after concat, over the full [txt+img] sequence. + # image_rotary_emb shape: [txt_len+img_len, head_dim] — matches q/k dim=1. + if image_rotary_emb is not None: + q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1) + k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1) + + x = self.attn(q, k, v, num_replicated_prefix=txt_seq_len) + x = x.flatten(2, 3).to(q.dtype) + + encoder_out, hidden_out = x.split_with_sizes( + [txt_seq_len, x.shape[1] - txt_seq_len], dim=1 + ) + hidden_out, _ = self.to_out[0](hidden_out) + encoder_out, _ = self.to_add_out(encoder_out) + return hidden_out, encoder_out + + +class _LongCatSingleAttention(nn.Module): + """Single-stream attention for _SingleTransformerBlock. + + txt and img are already concatenated by the block before calling here. + No output projection — the block handles proj_out itself. + + TP: Q/K/V use ColumnParallelLinear with gather_output=False (head-sharded). + USPAttention receives [B, S, H_local, D] directly; no all-gather needed here. + proj_mlp/proj_out in _SingleTransformerBlock must also use gather_output=False + so the concat [attn_output, mlp_hidden_states] is uniformly sharded. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + bias: bool = True, + eps: float = 1e-6, + prefix: str = "", + ): + super().__init__() + tp_size = get_tp_world_size() + self.num_local_heads = num_attention_heads // tp_size + assert ( + num_attention_heads % tp_size == 0 + ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + self.head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + self.norm_q = RMSNorm(attention_head_dim, eps=eps) + self.norm_k = RMSNorm(attention_head_dim, eps=eps) + + self.to_q = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_q" + ) + self.to_k = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_k" + ) + self.to_v = ColumnParallelLinear( + dim, inner_dim, bias=bias, gather_output=False, prefix=f"{prefix}.to_v" + ) + + self.attn = USPAttention( + num_heads=self.num_local_heads, + head_size=attention_head_dim, + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + q, _ = self.to_q(hidden_states) + k, _ = self.to_k(hidden_states) + v, _ = self.to_v(hidden_states) + q = q.unflatten(-1, (self.num_local_heads, self.head_dim)) + k = k.unflatten(-1, (self.num_local_heads, self.head_dim)) + v = v.unflatten(-1, (self.num_local_heads, self.head_dim)) + + # SGLang fused QK-norm + q, k = apply_qk_norm(q, k, self.norm_q, self.norm_k, self.head_dim) + + # RoPE via diffusers (supports full head_dim rotation, sequence_dim=1) + if image_rotary_emb is not None: + q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1) + k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1) + + x = self.attn(q, k, v) + return x.flatten(2, 3).to(q.dtype) + + +# --------------------------------------------------------------------------- +# Transformer blocks +# --------------------------------------------------------------------------- + + +class _SingleTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + prefix: str = "", + ): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + self.norm = AdaLayerNormZeroSingle(dim) + # proj_mlp: ColumnParallelLinear with gather_output=False keeps output + # head-sharded, consistent with attn_output from _LongCatSingleAttention. + self.proj_mlp = ColumnParallelLinear( + dim, + self.mlp_hidden_dim, + bias=True, + gather_output=False, + prefix=f"{prefix}.proj_mlp", + ) + self.act_mlp = nn.GELU(approximate="tanh") + # proj_out: RowParallelLinear reduces sharded [attn | mlp] concat via + # all-reduce, matching Flux2SingleTransformerBlockAttention.to_out. + self.proj_out = RowParallelLinear( + dim + self.mlp_hidden_dim, + dim, + bias=True, + input_is_parallel=True, + prefix=f"{prefix}.proj_out", + ) + self.attn = _LongCatSingleAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + bias=True, + eps=1e-6, + prefix=f"{prefix}.attn", + ) + tp_size = get_tp_world_size() + if tp_size > 1: + self._patch_proj_out_weight_loader( + inner_dim=num_attention_heads * attention_head_dim, + mlp_dim=self.mlp_hidden_dim, + tp_size=tp_size, + ) + + def _patch_proj_out_weight_loader( + self, inner_dim: int, mlp_dim: int, tp_size: int + ) -> None: + # proj_out input is [attn_shard | mlp_shard] where the two shards come + # from non-contiguous column ranges in the checkpoint weight matrix. + # Default RowParallelLinear.weight_loader slices contiguously, which is + # wrong here; override it to pick the correct columns per rank. + proj_out = self.proj_out + tp_rank = proj_out.tp_rank + + def _loader(param, loaded_weight): + input_dim = getattr(param, "input_dim", None) + if input_dim is not None: + a = inner_dim // tp_size + m = mlp_dim // tp_size + attn_cols = loaded_weight.narrow(input_dim, tp_rank * a, a) + mlp_cols = loaded_weight.narrow(input_dim, inner_dim + tp_rank * m, m) + param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim)) + else: + param.data.copy_(loaded_weight) + + proj_out.weight_loader = _loader + proj_out.weight.weight_loader = _loader + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb=None, + **kwargs, + ): + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + residual = hidden_states + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states, _ = self.proj_mlp(norm_hidden_states) + mlp_hidden_states = self.act_mlp(mlp_hidden_states) + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + gate = gate.unsqueeze(1) + hidden_states, _ = self.proj_out(hidden_states) + hidden_states = gate * hidden_states + hidden_states = residual + hidden_states + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + encoder_hidden_states, hidden_states = ( + hidden_states[:, :text_seq_len], + hidden_states[:, text_seq_len:], + ) + return encoder_hidden_states, hidden_states + + +class _TransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + qk_norm_eps: float = 1e-6, + prefix: str = "", + ): + super().__init__() + self.norm1 = AdaLayerNormZero(dim) + self.norm1_context = AdaLayerNormZero(dim) + self.attn = _LongCatJointAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + bias=True, + eps=qk_norm_eps, + prefix=f"{prefix}.attn", + ) + # norm2/norm2_context use eps=1e-6 matching diffusers (not configurable) + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = _LongCatFFN(dim=dim, inner_dim=dim * 4, prefix=f"{prefix}.ff") + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_context = _LongCatFFN( + dim=dim, inner_dim=dim * 4, prefix=f"{prefix}.ff_context" + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb=None, + **kwargs, + ): + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, emb=temb + ) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( + self.norm1_context(encoder_hidden_states, emb=temb) + ) + + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = ( + norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ) + ff_output = self.ff(norm_hidden_states) + ff_output = gate_mlp.unsqueeze(1) * ff_output + hidden_states = hidden_states + ff_output + + context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = ( + norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + + c_shift_mlp[:, None] + ) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = ( + encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + ) + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +# --------------------------------------------------------------------------- +# Position embedding +# --------------------------------------------------------------------------- + + +class _LongCatPosEmbed(nn.Module): + def __init__(self, theta: int, axes_dim: List[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (freqs_cos, freqs_sin) each of shape [S, head_dim].""" + cos_out, sin_out = [], [] + pos = ids.float() + freqs_dtype = ( + torch.float32 if ids.device.type in ("mps", "npu") else torch.float64 + ) + for i, dim in enumerate(self.axes_dim): + cos, sin = get_1d_rotary_pos_embed( + dim, + pos[:, i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +# --------------------------------------------------------------------------- +# Timestep embedding +# --------------------------------------------------------------------------- + + +class _TimestepEmbeddings(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = Timesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + + def forward(self, timestep: torch.Tensor, hidden_dtype) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + return self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +class LongCatImageTransformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin): + """SGLang implementation of the LongCat-Image transformer. + + Accepts raw scheduler timesteps (in [0, 1000]) and feeds them directly to + the timestep embedder (no /1000). The diffusers pipeline divides timesteps + by 1000 before calling its transformer, which multiplies them back by 1000 + internally — so the embedder receives the same value either way. + """ + + _aliases = ["LongCatImageTransformer2DModel"] + + param_names_mapping = {} + _fsdp_shard_conditions = [] + _compile_conditions = [] + + def __init__(self, config, hf_config: dict = None, quant_config=None): + super().__init__(config=config, hf_config=hf_config or {}) + arch = config.arch_config + + patch_size = getattr(arch, "patch_size", 1) + in_channels = getattr(arch, "in_channels", 64) + num_layers = getattr(arch, "num_layers", 19) + num_single_layers = getattr(arch, "num_single_layers", 38) + attention_head_dim = getattr(arch, "attention_head_dim", 128) + num_attention_heads = getattr(arch, "num_attention_heads", 24) + joint_attention_dim = getattr(arch, "joint_attention_dim", 3584) + axes_dims_rope = getattr(arch, "axes_dims_rope", [16, 56, 56]) + + self.config = config + self.out_channels = in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + # Required by BaseDiT.__post_init__ + self.hidden_size = self.inner_dim + self.num_attention_heads = num_attention_heads + self.num_channels_latents = 16 # unpacked latent channels + + self.pos_embed = _LongCatPosEmbed(theta=10000, axes_dim=axes_dims_rope) + self.time_embed = _TimestepEmbeddings(embedding_dim=self.inner_dim) + self.context_embedder = ColumnParallelLinear( + joint_attention_dim, + self.inner_dim, + bias=True, + gather_output=True, + prefix="context_embedder", + ) + self.x_embedder = ColumnParallelLinear( + in_channels, + self.inner_dim, + bias=True, + gather_output=True, + prefix="x_embedder", + ) + + self.transformer_blocks = nn.ModuleList( + [ + _TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + ) + + self.single_transformer_blocks = nn.ModuleList( + [ + _SingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"single_transformer_blocks.{i}", + ) + for i in range(num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 + ) + self.proj_out = ColumnParallelLinear( + self.inner_dim, + patch_size * patch_size * self.out_channels, + bias=True, + gather_output=True, + prefix="proj_out", + ) + + self.layer_names = ["transformer_blocks", "single_transformer_blocks"] + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.Tensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor = None, + return_dict: bool = False, + **kwargs, + ): + hidden_states, _ = self.x_embedder(hidden_states) + + # Raw scheduler timestep in [0, 1000] — fed directly (see module docstring). + temb = self.time_embed(timestep.to(hidden_states.dtype), hidden_states.dtype) + + encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states) + + # RoPE is computed on the fly from txt_ids + img_ids, matching diffusers' + # transformer (which calls self.pos_embed internally). image_rotary_emb is + # only passed in for standalone model testing. + image_rotary_emb = kwargs.get("image_rotary_emb") or self.pos_embed( + torch.cat((txt_ids, img_ids), dim=0) + ) + + for block in self.transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + for block in self.single_transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = self.norm_out(hidden_states, temb) + output, _ = self.proj_out(hidden_states) + + if return_dict: + from diffusers.models.modeling_outputs import Transformer2DModelOutput + + return Transformer2DModelOutput(sample=output) + return output + + +EntryClass = LongCatImageTransformer2DModel diff --git a/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py b/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py new file mode 100644 index 000000000..dd0a31067 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py @@ -0,0 +1,80 @@ +"""LongCat-Image pipeline 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.model_specific_stages.longcat_image import ( + LongCatPromptRewriteStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + + +def _prepare_mu(batch, server_args): + """Compute mu for FlowMatchEulerDiscreteScheduler from the packed latent token count.""" + from sglang.multimodal_gen.configs.pipeline_configs.longcat_image import ( + _calculate_shift, + ) + + image_seq_len = batch.latents.shape[1] + mu = _calculate_shift(image_seq_len) + return "mu", mu + + +class LongCatImagePipeline(LoRAPipeline, ComposedPipelineBase): + """Pipeline for LongCat-Image text-to-image generation.""" + + pipeline_name = "LongCatImagePipeline" + + # The Qwen2.5-VL text encoder is loaded in-stage by LongCatPromptRewriteStage + # (not via TextEncoderLoader), so "text_encoder" is intentionally absent; + # the stage registers the loaded module via add_module("text_encoder", ...) + # so the standard TextEncodingStage can fetch the same instance. + _required_config_modules = [ + "tokenizer", + "text_processor", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + # 1. Prompt rewriting (optional) + request-level setup (generator, cfg renorm). + # Loads the HF Qwen2.5-VL encoder and shares it with TextEncodingStage. + rewrite_stage = LongCatPromptRewriteStage( + tokenizer=self.get_module("tokenizer"), + text_processor=self.get_module("text_processor"), + model_path=self.model_path, + text_encoder_dtype=PRECISION_TO_TYPE[ + server_args.pipeline_config.text_encoder_precisions[0] + ], + ) + self.add_stage(rewrite_stage) + self.add_module("text_encoder", rewrite_stage.text_encoder) + + # 2. Text encoding via the standard stage (tokenize_prompt + + # postprocess_text_funcs hooks on the pipeline config). Shares the + # encoder instance registered above; both stages declare a + # "text_encoder" ComponentUse so the residency manager keeps it + # resident across rewrite->encode and offloads after the last use. + self.add_standard_text_encoding_stage() + + # 3. Latent preparation (batch-size-aware via pipeline config hooks) + self.add_standard_latent_preparation_stage() + + # 4. Timestep preparation (mu computed from packed latent token count) + self.add_standard_timestep_preparation_stage( + prepare_extra_kwargs=[_prepare_mu], + ) + + # 5. Standard denoising loop. txt_ids/img_ids are built per-step inside + # prepare_*_cond_kwargs; the DiT computes RoPE on the fly from them + # (matching diffusers' transformer). + self.add_standard_denoising_stage() + + # 6. Standard VAE decoding + self.add_standard_decoding_stage() + + +EntryClass = [LongCatImagePipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image.py new file mode 100644 index 000000000..5ea41183b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image.py @@ -0,0 +1,322 @@ +"""Prompt-rewriting stage for LongCat-Image (T2I). + +`LongCatPromptRewriteStage` optionally rewrites the prompt via the HuggingFace +`Qwen2_5_VLForConditionalGeneration.generate()` and sets the CPU generator for +seed reproducibility. Text encoding, latent preparation, RoPE and denoising are +all handled by the standard stages + `LongCatImagePipelineConfig` hooks, so this +is the only model-specific stage. +""" + +import re +from typing import List + +import torch + +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__) + +# Budget for the prompt-rewrite AR generation. Matches the official diffusers +# LongCatImagePipeline (which uses tokenizer_max_length=512). +REWRITE_MAX_NEW_TOKENS = 512 + +# Prompt-rewrite system messages (one per language). Copied and adapted from +# https://github.com/huggingface/diffusers (src/diffusers/pipelines/longcat_image/system_messages.py). +SYSTEM_PROMPT_EN = """ +You are a prompt engineering expert for text-to-image models. Since text-to-image models have limited capabilities in +understanding user prompts, you need to identify the core theme and intent of the user's input and improve the model's +understanding accuracy and generation quality through optimization and rewriting. The rewrite must strictly retain all +information from the user's original prompt without deleting or distorting any details. Specific requirements are as +follows: +1. The rewrite must not affect any information expressed in the user's original prompt; the rewritten prompt should use + coherent natural language, avoid low-information redundant descriptions, and keep the rewritten prompt length as + concise as possible. +2. Ensure consistency between input and output languages: Chinese input yields Chinese output, and English input yields + English output. The rewritten token count should not exceed 512. +3. The rewritten description should further refine subject characteristics and aesthetic techniques appearing in the + original prompt, such as lighting and textures. +4. If the original prompt does not specify an image style, ensure the rewritten prompt uses a **realistic photography + style**. If the user specifies a style, retain the user's style. +5. When the original prompt requires reasoning to clarify user intent, use logical reasoning based on world knowledge + to convert vague abstract descriptions into specific tangible objects (e.g., convert "the tallest animal" to "a + giraffe"). +6. When the original prompt requires text generation, please use double quotes to enclose the text part (e.g., `"50% + OFF"`). +7. When the original prompt requires generating text-heavy scenes like webpages, logos, UIs, or posters, and no + specific text content is specified, you need to infer appropriate text content and enclose it in double quotes. For + example, if the user inputs: "A tourism flyer with a grassland theme," it should be rewritten as: "A tourism flyer + with the image title 'Grassland'." +8. When negative words exist in the original prompt, ensure the rewritten prompt does not contain negative words. For + example, "a lakeside without boats" should be rewritten such that the word "boat" does not appear at all. +9. Except for text content explicitly requested by the user, **adding any extra text content is prohibited**. +Here are examples of rewrites for different types of prompts: # Examples (Few-Shot Learning) + 1. User Input: An animal with nine lives. + Rewrite Output: A cat bathed in soft sunlight, its fur soft and glossy. The background is a comfortable home + environment with light from the window filtering through curtains, creating a warm light and shadow effect. The + shot uses a medium distance perspective to highlight the cat's leisurely and stretched posture. Light cleverly hits + the cat's face, emphasizing its spirited eyes and delicate whiskers, adding depth and affinity to the image. + 2. User Input: Create an anime-style tourism flyer with a grassland theme. + Rewrite Output: In the lower right of the center, a short-haired girl sits sideways on a gray, irregularly shaped + rock. She wears a white short-sleeved dress and brown flat shoes, holding a bunch of small white flowers in her + left hand, smiling with her legs hanging naturally. The girl has dark brown shoulder-length hair with bangs + covering her forehead, brown eyes, and a slightly open mouth. The rock surface has textures of varying depths. To + the girl's left and front is lush grass, with long, yellow-green blades, some glowing golden in the sunlight. The + grass extends into the distance, forming rolling green hills that fade in color as they recede. The sky occupies + the upper half of the picture, pale blue dotted with a few fluffy white clouds. In the upper left corner, there is + a line of text in italic, dark green font reading "Explore Nature's Peace". Colors are dominated by green, blue, + and yellow, fluid lines, and distinct light and shadow contrast, creating a quiet and comfortable atmosphere. + 3. User Input: A Christmas sale poster with a red background, promoting a Buy 1 Get 1 Free milk tea offer. + Rewrite Output: The poster features an overall red tone, embellished with white snowflake patterns on the top and + left side. The upper right features a bunch of holly leaves with red berries and a pine cone. In the upper center, + golden 3D text reads "Christmas Heartwarming Feedback" centered, along with red bold text "Buy 1 Get 1". Below, two + transparent cups filled with bubble tea are placed side by side; the tea is light brown with dark brown pearls + scattered at the bottom and middle. Below the cups, white snow piles up, decorated with pine branches, red berries, + and pine cones. A blurry Christmas tree is faintly visible in the lower right corner. The image has high clarity, + accurate text content, a unified design style, a prominent Christmas theme, and a reasonable layout, providing + strong visual appeal. + 4. User Input: A woman indoors shot in natural light, smiling with arms crossed, showing a relaxed and confident + posture. + Rewrite Output: The image features a young Asian woman with long dark brown hair naturally falling over her + shoulders, with some strands illuminated by light, showing a soft sheen. Her features are delicate, with long + eyebrows, bright and spirited dark brown eyes looking directly at the camera, revealing peace and confidence. She + has a high nose bridge, full lips with nude lipstick, and corners of the mouth slightly raised in a faint smile. + Her skin is fair, with cheeks and collarbones illuminated by warm light, showing a healthy ruddiness. She wears a + black spaghetti strap tank top revealing graceful collarbone lines, and a thin gold necklace with small beads and + metal bars glinting in the light. Her outer layer is a beige knitted cardigan, soft in texture with visible + knitting patterns on the sleeves. Her arms are crossed over her chest, hands covered by the cardigan sleeves, in a + relaxed posture. The background is a pure dark brown without extra decoration, making the figure the absolute + focus. The figure is located in the center of the frame. Light enters from the upper right, creating bright spots + on her left cheek, neck, and collarbone, while the right side is slightly shadowed, creating a three-dimensional + and soft tone. Image details are clear, showcasing skin texture, hair, and clothing materials well. Colors are + dominated by warm tones, with the combination of beige and dark brown creating a warm and comfortable atmosphere. + The overall style is natural, elegant, and artistic. + 5. User Input: Create a series of images showing the growth process of an apple from seed to fruit. The series should + include four stages: 1. Sowing, 2. Seedling growth, 3. Plant maturity, 4. Fruit harvesting. + Rewrite Output: A 4-panel exquisite illustration depicting the growth process of an apple, capturing each stage + precisely and clearly. 1. "Sowing": A close-up shot of a hand gently placing a small apple seed into fertile dark + soil, with visible soil texture and the seed's smooth surface. The background is a soft-focus garden dotted with + green leaves and sunlight filtering through. 2. "Seedling Growth": A young apple sapling breaks through the soil, + stretching tender green leaves toward the sky. The scene is set in a vibrant garden illuminated by warm golden + light, highlighting the seedling's delicate structure. 3. "Plant Maturity": A mature apple tree, lush with branches + and leaves, covered in tender green foliage and developing small apples. The background is a vibrant orchard under + a clear blue sky, with dappled sunlight creating a peaceful atmosphere. 4. "Fruit Harvesting": A hand reaches into + the tree to pick a ripe red apple, its smooth skin glistening in the sun. The scene shows the abundance of the + orchard, with baskets of apples in the background, giving a sense of fulfillment. Each illustration uses a + realistic style, focusing on details and harmonious colors to showcase the natural beauty and development of the + apple's life cycle. + 6. User Input: If 1 represents red, 2 represents green, 3 represents purple, and 4 represents yellow, please generate + a four-color rainbow based on this rule. The color order from top to bottom is 3142. + Rewrite Output: The image consists of four horizontally arranged colored stripes, ordered from top to bottom as + purple, red, yellow, and green. A white number is centered on each stripe. The top purple stripe features the + number "3", the red stripe below it has the number "1", the yellow stripe further down has the number "4", and the + bottom green stripe has the number "2". All numbers use a sans-serif font in pure white, forming a sharp contrast + with the background colors to ensure good readability. The stripes have high color saturation and a slight texture. + The overall layout is simple and clear, with distinct visual effects and no extra decorative elements, emphasizing + the numerical information. The image is high definition, with accurate colors and a consistent style, offering + strong visual appeal. + 7. User Input: A stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", natural light, background is a + Chinese garden. + Rewrite Output: An ancient stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", the surface covered with + traces of time, the writing clear and deep. Natural light falls from above, softly illuminating every detail of the + stone tablet and enhancing its sense of history. The background is an elegant Chinese garden featuring lush bamboo + forests, winding paths, and quiet pools, creating a serene and distant atmosphere. The overall picture uses a + realistic style with rich details and natural light and shadow effects, highlighting the cultural heritage of the + stone tablet and the classical beauty of the garden. +# Output Format Please directly output the rewritten and optimized Prompt content. Do not include any explanatory +language or JSON formatting, and do not add opening or closing quotes yourself.""" + +SYSTEM_PROMPT_ZH = """ +你是一名文生图模型的prompt +engineering专家。由于文生图模型对用户prompt的理解能力有限,你需要识别用户输入的核心主题和意图,并通过优化改写提升模型的理解准确性和生成质量。改写必须严格保留用户原始prompt的所有信息,不得删减或曲解任何细节。 +具体要求如下: +1. 改写不能影响用户原始prompt里表达的任何信息,改写后的prompt应该使用连贯的自然语言表达,不要出现低信息量的冗余描述,尽可能保持改写后prompt长度精简。 +2. 请确保输入和输出的语言类型一致,中文输入中文输出,英文输入英文输出,改写后的token数量不要超过512个; +3. 改写后的描述应当进一步完善原始prompt中出现的主体特征、美学技巧,如打光、纹理等; +4. 如果原始prompt没有指定图片风格时,确保改写后的prompt使用真实摄影风格,如果用户指定了图片风格,则保留用户风格; +5. 当原始prompt需要推理才能明确用户意图时,根据世界知识进行适当逻辑推理,将模糊抽象描述转化为具体指向事物(例:将"最高的动物"转化为"一头长颈鹿")。 +6. 当原始prompt需要生成文字时,请使用双引号圈定文字部分,例:`"限时5折"`)。 +7. 当原始prompt需要生成网页、logo、ui、海报等文字场景时,且没有指定具体的文字内容时,需要推断出合适的文字内容,并使用双引号圈定,如用户输入:一个旅游宣传单,以草原为主题。应该改写成:一个旅游宣传单,图片标题为"草原"。 +8. 当原始prompt中存在否定词时,需要确保改写后的prompt不存在否定词,如没有船的湖边,改写后的prompt不能出现船这个词汇。 +9. 除非用户指定生成品牌logo,否则不要增加额外的品牌logo. +10. 除了用户明确要求书写的文字内容外,**禁止增加任何额外的文字内容**。 +以下是针对不同类型prompt改写的示例: + +# Examples (Few-Shot Learning) + 1. 用户输入: 九条命的动物。 + 改写输出: + 一只猫,被柔和的阳光笼罩着,毛发柔软而富有光泽。背景是一个舒适的家居环境,窗外的光线透过窗帘,形成温馨的光影效果。镜头采用中距离视角,突出猫悠闲舒展的姿态。光线巧妙地打在猫的脸部,强调它灵动的眼睛和精致的胡须,增加画面的层次感与亲和力。 + 2. 用户输入: 制作一个动画风格的旅游宣传单,以草原为主题。 + 改写输出: + 画面中央偏右下角,一个短发女孩侧身坐在灰色的不规则形状岩石上,她穿着白色短袖连衣裙和棕色平底鞋,左手拿着一束白色小花,面带微笑,双腿自然垂下。女孩的头发为深棕色,齐肩短发,刘海覆盖额头,眼睛呈棕色,嘴巴微张。岩石表面有深浅不一的纹理。女孩的左侧和前方是茂盛的草地,草叶细长,呈黄绿色,部分草叶在阳光下泛着金色的光芒,仿佛被阳光照亮。草地向远处延伸,形成连绵起伏的绿色山丘,山丘的颜色由近及远逐渐变浅。天空占据了画面的上半部分,呈淡蓝色,点缀着几朵白色蓬松的云彩。画面的左上角有一行文字,文字内容是斜体、深绿色的"Explore + Nature's Peace"。色彩以绿色、蓝色和黄色为主,线条流畅,光影明暗对比明显,营造出一种宁静、舒适的氛围。 + 3. 用户输入: 一张以红色为背景的圣诞节促销海报,主要宣传奶茶买一送一的优惠活动。 + 改写输出: 海报整体呈现红色调,上方和左侧点缀着白色雪花图案,右上方有一束冬青叶和红色浆果,以及一个松果。海报中央偏上位置,金色立体字样"圣诞节 + 暖心回馈"居中排列,和红色粗体字"买1送1"。海报下方,两个装满珍珠奶茶的透明杯子并排摆放,杯中奶茶呈浅棕色,底部和中间散布着深棕色珍珠。杯子下方,堆积着白色雪花,雪花上装饰着松枝、红色浆果和松果。右下角隐约可见一棵模糊的圣诞树。图片清晰度高,文字内容准确,整体设计风格统一,圣诞主题突出,排版布局合理,具有较强的视觉吸引力。 + 4. 用户输入: 一位女性在室内以自然光线拍摄,她面带微笑,双臂交叉,展现出轻松自信的姿态。 + 改写输出: + 画面中是一位年轻的亚洲女性,她拥有深棕色的长发,发丝自然地垂落在双肩,部分发丝被光线照亮,呈现出柔和的光泽。她的五官精致,眉毛修长,眼睛明亮有神,瞳孔呈深棕色,眼神直视镜头,流露出平和与自信。鼻梁挺拔,嘴唇丰满,涂有裸色系唇膏,嘴角微微上扬,展现出浅浅的微笑。她的肤色白皙,脸颊和锁骨处被暖色调的光线照亮,呈现出健康的红润感。她穿着一件黑色的细吊带背心,肩带纤细,露出优美的锁骨线条。脖颈上佩戴着一条金色的细项链,项链由小珠子和几个细长的金属条组成,在光线下闪烁着光泽。她的外搭是一件米黄色的针织开衫,材质柔软,袖子部分有明显的针织纹理。她双臂交叉在胸前,双手被开衫的袖子覆盖,姿态放松。背景是纯粹的深棕色,没有多余的装饰,使得人物成为画面的绝对焦点。人物位于画面中央。光线从画面的右上方射入,在人物的左侧脸颊、脖颈和锁骨处形成明亮的光斑,右侧则略显阴影,营造出立体感和柔和的影调。图像细节清晰,人物的皮肤纹理、发丝以及衣物材质都得到了很好的展现。色彩以暖色调为主,米黄色和深棕色的搭配营造出温馨舒适的氛围。整体呈现出一种自然、优雅且富有亲和力的艺术风格。 + 5. 用户输入:创作一系列图片,展现苹果从种子到结果的生长过程。该系列图片应包含以下四个阶段:1. 播种,2. 幼苗生长,3. 植物成熟,4. 果实采摘。 + 改写输出:一个4宫格的精美插图,描绘苹果的生长过程,精确清晰地捕捉每个阶段。1."播种":特写镜头,一只手轻轻地将一颗小小的苹果种子放入肥沃的深色土壤中,土壤的纹理和种子光滑的表面清晰可见。背景是花园的柔焦画面,点缀着绿色的树叶和透过树叶洒下的阳光。2."幼苗生长":一棵幼小的苹果树苗破土而出,嫩绿的叶子向天空舒展。场景设定在一个生机勃勃的花园中,温暖的金光照亮了它。幼苗的纤细结构。3."植物的成熟":一棵成熟的苹果树,枝繁叶茂,挂满了嫩绿的叶子和正在萌发的小苹果。背景是一片生机勃勃的果园,湛蓝的天空下,斑驳的阳光营造出宁静祥和的氛围。4."采摘果实":一只手伸向树上,摘下一个成熟的红苹果,苹果光滑的果皮在阳光下闪闪发光。画面展现了果园的丰收景象,背景中摆放着一篮篮的苹果,给人一种圆满满足的感觉。每幅插图都采用写实风格,注重细节,色彩和谐,展现了苹果生命周期的自然之美和发展过程。 + 6. 用户输入: 如果1代表红色,2代表绿色,3代表紫色,4代表黄色,请按照此规则生成四色彩虹。它的颜色顺序从上到下是3142 + 改写输出:图片由四个水平排列的彩色条纹组成,从上到下依次为紫色、红色、黄色和绿色。每个条纹上都居中放置一个白色数字。最上方的紫色条纹上是数字"3",其下方红色条纹上是数字"1",再下方黄色条纹上是数字"4",最下方的绿色条纹上是数字"2"。所有数字均采用无衬线字体,颜色为纯白色,与背景色形成鲜明对比,确保了良好的可读性。条纹的颜色饱和度高,且带有轻微的纹理感,整体排版简洁明了,视觉效果清晰,没有多余的装饰元素,强调了数字信息本身。图片整体清晰度高,色彩准确,风格一致,具有较强的视觉吸引力。 + 7. 用户输入:石碑上刻着"关关雎鸠,在河之洲",自然光照,背景是中式园林 + 改写输出:一块古老的石碑上刻着"关关雎鸠,在河之洲",石碑表面布满岁月的痕迹,字迹清晰而深刻。自然光线从上方洒下,柔和地照亮石碑的每一个细节,增强了其历史感。背景是一座典雅的中式园林,园林中有翠绿的竹林、蜿蜒的小径和静谧的水池,营造出一种宁静而悠远的氛围。整体画面采用写实风格,细节丰富,光影效果自然,突出了石碑的文化底蕴和园林的古典美。 +# 输出格式 请直接输出改写优化后的 Prompt 内容,不要包含任何解释性语言或 JSON 格式,不要自行添加开头或结尾的引号。 +""" + + +def _get_prompt_language(prompt): + pattern = re.compile(r"[一-鿿]") + return "zh" if bool(pattern.search(prompt)) else "en" + + +class LongCatPromptRewriteStage(PipelineStage): + """Optional prompt rewriting + request-level setup for LongCat-Image. + + Loads the Qwen2.5-VL text encoder (HuggingFace) in-stage and, when + `enable_prompt_rewrite` is set, rewrites the prompt via `.generate()` + (using the checkpoint's generation_config.json sampling params). Always + sets the CPU generator for seed reproducibility. CFG-renorm params are + read directly from sampling_params in `postprocess_cfg_noise`, not set here. + + The same encoder instance is shared with the standard `TextEncodingStage` + (the pipeline registers it via `add_module("text_encoder", ...)`), so + rewrite and encode run on one set of weights. Both stages declare a + `text_encoder` ComponentUse; the residency manager keeps the encoder + resident across the adjacent rewrite->encode uses and offloads it only + after the last use (when `--text-encoder-cpu-offload` is enabled). + """ + + def __init__( + self, + tokenizer, + text_processor, + model_path: str, + text_encoder_dtype: torch.dtype, + ): + super().__init__() + self.text_encoder_dtype = text_encoder_dtype + from transformers import Qwen2_5_VLForConditionalGeneration + + cpu_offload = self.server_args.text_encoder_cpu_offload + init_device = torch.device("cpu") if cpu_offload else get_local_torch_device() + self.text_encoder = ( + Qwen2_5_VLForConditionalGeneration.from_pretrained( + model_path, subfolder="text_encoder" + ) + .to(init_device) + .to(dtype=self.text_encoder_dtype) + ) + self.tokenizer = tokenizer + self.text_processor = text_processor + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + # "text_encoder" matches is_text_encoder_component_name, so + # text_encoder_cpu_offload routes through VanillaD2HStrategy (the encoder + # is a plain HF nn.Module, not FSDP-sharded). memory_intensive triggers + # torch.cuda.empty_cache() after the encoder leaves CUDA. + return [ + ComponentUse( + stage_name, + "text_encoder", + target_dtype=self.text_encoder_dtype, + memory_intensive=True, + ), + ] + + def _rewire_prompt(self, prompt: List[str], device: torch.device) -> List[str]: + """Rewrite prompts via `.generate()` on the Qwen2.5-VL encoder. + + Mirrors diffusers LongCatImagePipeline.rewire_prompt(): only + `max_new_tokens` is passed, so `.generate()` uses the checkpoint's + `generation_config.json` (top_k=1, top_p=0.001, temperature=0.1, + repetition_penalty=1.05). Do not force do_sample/num_beams — that + diverges from the reference and changes the rewritten prompt. + """ + all_text = [] + for each_prompt in prompt: + language = _get_prompt_language(each_prompt) + if language == "zh": + question = ( + SYSTEM_PROMPT_ZH + + f"\n用户输入为:{each_prompt}\n改写后的prompt为:" + ) + else: + question = ( + SYSTEM_PROMPT_EN + f"\nUser Input: {each_prompt}\nRewritten prompt:" + ) + message = [ + { + "role": "user", + "content": [{"type": "text", "text": question}], + } + ] + text = self.text_processor.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) + all_text.append(text) + + inputs = self.text_processor( + text=all_text, padding=True, return_tensors="pt" + ).to(device) + + with set_forward_context(current_timestep=0, attn_metadata=None): + # Match the reference: only max_new_tokens is passed; the checkpoint's + # generation_config.json supplies the sampling params. + generated_ids = self.text_encoder.generate( + **inputs, + max_new_tokens=REWRITE_MAX_NEW_TOKENS, + ) + + prompt_len = inputs["input_ids"].shape[1] + generated_ids_trimmed = generated_ids[:, prompt_len:] + rewritten = self.text_processor.batch_decode( + generated_ids_trimmed, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) + logger.info("Rewritten prompts: %s", rewritten) + return rewritten + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = get_local_torch_device() + + enable_prompt_rewrite = getattr(batch, "enable_prompt_rewrite", True) + + 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 + + if enable_prompt_rewrite: + logger.info( + "Prompt rewriting is enabled (enable_prompt_rewrite=True). " + "This runs autoregressive decoding on the Qwen2.5-VL text " + "encoder (up to %d tokens). Pass --enable-prompt-rewrite " + "false to skip.", + REWRITE_MAX_NEW_TOKENS, + ) + prompt = batch.prompt + if isinstance(prompt, str): + prompt = [prompt] + batch.prompt = self._rewire_prompt(prompt, device) + + # Always set: the CPU generator governs latent noise for seed + # reproducibility (diffusers randn_tensor generates on CPU then moves + # to device, so CPU and CUDA generators differ for the same seed). + batch.generator = torch.Generator(device="cpu").manual_seed(batch.seed) + return batch