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

[diffusion] model: support LongCat-Image (#23274)

Browse Source 
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Liu Zhenlong
2026-08-13 19:48:26 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent eea2e5d6e5
commit b764194e81
10 changed files with 1641 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/longcat_image.py
+28
View File
@@ -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"
python/sglang/multimodal_gen/configs/models/vaes/longcat_image.py
+24
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/longcat_image.py
+422
View File
@@ -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
python/sglang/multimodal_gen/configs/sample/longcat_image.py
+14
View File
@@ -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
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+22
View File
@@ -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",
python/sglang/multimodal_gen/registry.py
+18
View File
@@ -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()
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+1 -1
View File
@@ -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(
python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
+710
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines/longcat_image.py
+80
View File
@@ -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]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image.py
+322
View File
@@ -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