[Feature] Support Qwen/Qwen-Image-2512

[pengchengneo]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

as title, please pasted your detailed RFC for this model before dive into it, which can make sure your idea is aligned with our multimodal framework

Related resources

https://huggingface.co/Qwen/Qwen-Image-2512

[andy1126]

+1

[andy1126]

[RFC] Support Qwen-Image-2512 in SGL_JAX

Objective: Add Qwen-Image-2512 text-to-image / image-to-image generation model support to the sglang-jax inference framework.

What This PR Would Be Adding

New Files

File	Description
python/sgl_jax/srt/multimodal/configs/dits/qwen_image_config.py	DiT configuration
python/sgl_jax/srt/multimodal/configs/vae/qwen_image_vae_config.py	VAE configuration
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_generation.py	Model Entrypoint
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_dit_block.py	Transformer block
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_vae.py	VAE encoder/decoder
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_text_encoder.py	Text encoder wrapper
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_vision_encoder.py	Vision encoder
python/sgl_jax/srt/multimodal/models/qwen_image/qwen_image_weights_mapping.py	Weight mappings
python/sgl_jax/srt/multimodal/models/static_configs/qwen_image_t2i_stage_config.yaml	T2I pipeline config
python/sgl_jax/srt/multimodal/models/static_configs/qwen_image_i2i_stage_config.yaml	I2I pipeline config

Modified Files

File	Description
python/sgl_jax/srt/multimodal/manager/stage.py	add Qwen Image to get_model_class
python/sgl_jax/srt/multimodal/configs/config_registry.py	add Qwen Image to ConfigRegistry
python/sgl_jax/srt/multimodal/models/static_configs/yaml_registry.py	add Qwen Image pipeline configs to YamlRegistry
python/sgl_jax/srt/multimodal/model_executor/diffusion/diffusion_model_runner.py	Add 4D latent matrix for image or unsqueezing the shape to make it [B,C,T,H,W]
python/sgl_jax/srt/multimodal/manager/scheduler/vae_scheduler.py	Add support to process 2D images, for no temporal padding is needed

Implementation Details

1. DiT Configuration

File: configs/dits/qwen_image_config.py

Design Rationale:

• Aligned with WanModelConfig structure for consistency
• Removes 3D-related parameters (temporal dimension)
• Adds MMDiT dual-stream specific fields (dual_stream_depth, joint_depth)
• Includes I2I-specific optional fields (image_dim, added_kv_proj_dim)

2. Model Core Implementation

File: models/qwen_image/qwen_image_generation.py

Design Rationale:

• Dual-stream architecture: Text and image processed independently in early layers, then jointly in subsquent layers. This allows modality-specific processing before cross-modal fusion.
• AdaLN-Zero conditioning: Timestep condition injected via adaptive layer normalization with zero-initialized gates, following DiT paper for training stability.
• 2D position encoding: Uses 2D RoPE instead of 3D, matching image spatial structure.

3. Transformer Block

File: models/qwen_image/qwen_image_dit_block.py

Design Rationale:

• AdaLN-Zero: Generates (shift, scale, gate) from timestep for both attention and MLP. Zero-initialized gates ensure identity mapping at initialization.
• QK Norm: Applies RMSNorm to Q and K for training stability.
• I2I cross-attention: When image_context is provided, computes additional attention with reference image features.

4. VAE Implementation

File: models/qwen_image/qwen_image_vae.py

Design Rationale:

• Simplified from wanvae.py by removing temporal (3D) processing,
• Contains both encode (for I2I) and decode (for output) methods

5. Text Encoder Wrapper

File: models/qwen_image/qwen_image_text_encoder.py

Design Rationale:

• Qwen-Image uses Qwen2.5-VL as text encoder
• Only need embedding + transformer layers, no LM head
• Outputs hidden states for DiT conditioning

6. Vision Encoder for I2I

File: models/qwen_image/qwen_image_vision_encoder.py

Design Rationale:

• Reuses Qwen2.5-VL's Vision Transformer for encoding reference images
• Outputs semantic features for I2I cross-attention conditioning

Data Flow

Text-to-Image Pipeline

Text Encoder -> DiT -> VAE Decoder

Image-to-Image Pipeline

Text Encoder ------------------------->
| -----> DiT -> VAE Decoder
Vision Encoder -> VAE Encoder ------->

Task Items

• T2I Pipeline
  • DiT
  • VAE
  • Text Encoder
• I2I Pipeline
  • Vision Encoder
  • I2I Cross-Attention
• Full Testing

References

Qwen-Image-2512 HuggingFace