From b564132e44f5ec7b05f2dd5d36c320caf17b4bff Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 21 Jan 2026 16:53:18 +0800 Subject: [PATCH] [Diffusion][Acceleration] Support TeaCache for Z-Image (#817) Signed-off-by: gcanlin Co-authored-by: Hongsheng Liu --- docs/user_guide/diffusion_acceleration.md | 2 +- vllm_omni/diffusion/cache/teacache/config.py | 9 + .../diffusion/cache/teacache/extractors.py | 187 ++++++++++++++++++ .../models/z_image/z_image_transformer.py | 3 +- 4 files changed, 199 insertions(+), 2 deletions(-) diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 0184a8fcd..4572b7115 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -47,7 +47,7 @@ The following table shows which models are currently supported by each accelerat | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | ❌ | ❌ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ | | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ | diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 916832bbf..bb5e7cb27 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -25,6 +25,15 @@ _MODEL_COEFFICIENTS = { 3.20000000e00, -2.00000000e-02, ], + # Z-Image transformer coefficients + # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future + "ZImageTransformer2DModel": [ + -4.50000000e02, + 2.80000000e02, + -4.50000000e01, + 3.20000000e00, + -2.00000000e-02, + ], } diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index d3002ba57..5c35e0c8c 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -254,6 +254,192 @@ def extract_qwen_context( ) +def extract_zimage_context( + module: nn.Module, + x: list[torch.Tensor], + t: torch.Tensor, + cap_feats: list[torch.Tensor], + patch_size: int = 2, + f_patch_size: int = 1, + **kwargs: Any, +) -> CacheContext: + """ + Extract cache context for ZImageTransformer2DModel. + + This is the ONLY Z-Image-specific code needed for TeaCache support. + It encapsulates preprocessing, modulated input extraction, transformer execution, + and postprocessing logic. + + Args: + module: ZImageTransformer2DModel instance + x: List of image tensors per batch item + t: Timestep tensor + cap_feats: List of caption feature tensors per batch item + patch_size: Patch size for patchification (default: 2) + f_patch_size: Frame patch size (default: 1) + **kwargs: Additional keyword arguments ignored by this extractor + + Returns: + CacheContext with all information needed for generic caching + """ + from torch.nn.utils.rnn import pad_sequence + + if not hasattr(module, "layers") or len(module.layers) == 0: + raise ValueError("Module must have main transformer layers") + + bsz = len(x) + device = x[0].device + + # ============================================================================ + # PREPROCESSING (Z-Image specific) + # ============================================================================ + # Scale timestep and create timestep embedding + t_scaled = t * module.t_scale + adaln_input = module.t_embedder(t_scaled) + + # Patchify and embed inputs + ( + x_patches, + cap_feats_processed, + x_size, + x_pos_ids, + cap_pos_ids, + x_inner_pad_mask, + cap_inner_pad_mask, + ) = module.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + + # Process image patches through embedder and noise refiner + x_item_seqlens = [len(_) for _ in x_patches] + x_max_item_seqlen = max(x_item_seqlens) + + x_embedded = torch.cat(x_patches, dim=0) + x_embedded = module.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_embedded) + + # Match adaln_input dtype to x_embedded + adaln_input = adaln_input.type_as(x_embedded) + + # Apply pad token + x_embedded[torch.cat(x_inner_pad_mask)] = module.x_pad_token + x_list = list(x_embedded.split(x_item_seqlens, dim=0)) + + # Compute rope embeddings for image patches + x_cos, x_sin = module.rope_embedder(torch.cat(x_pos_ids, dim=0)) + x_cos = list(x_cos.split(x_item_seqlens, dim=0)) + x_sin = list(x_sin.split(x_item_seqlens, dim=0)) + + # Pad sequences for batch processing + x_batched = pad_sequence(x_list, batch_first=True, padding_value=0.0) + x_cos_batched = pad_sequence(x_cos, batch_first=True, padding_value=0.0) + x_sin_batched = pad_sequence(x_sin, batch_first=True, padding_value=0.0) + x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(x_item_seqlens): + x_attn_mask[i, :seq_len] = 1 + + # Run noise refiner blocks + for layer in module.noise_refiner: + x_batched = layer(x_batched, x_attn_mask, x_cos_batched, x_sin_batched, adaln_input) + + # Process caption features through embedder and context refiner + cap_item_seqlens = [len(_) for _ in cap_feats_processed] + cap_max_item_seqlen = max(cap_item_seqlens) + + cap_embedded = torch.cat(cap_feats_processed, dim=0) + cap_embedded = module.cap_embedder(cap_embedded) + cap_embedded[torch.cat(cap_inner_pad_mask)] = module.cap_pad_token + cap_list = list(cap_embedded.split(cap_item_seqlens, dim=0)) + + # Compute rope embeddings for caption + cap_cos, cap_sin = module.rope_embedder(torch.cat(cap_pos_ids, dim=0)) + cap_cos = list(cap_cos.split(cap_item_seqlens, dim=0)) + cap_sin = list(cap_sin.split(cap_item_seqlens, dim=0)) + + # Pad sequences for batch processing + cap_batched = pad_sequence(cap_list, batch_first=True, padding_value=0.0) + cap_cos_batched = pad_sequence(cap_cos, batch_first=True, padding_value=0.0) + cap_sin_batched = pad_sequence(cap_sin, batch_first=True, padding_value=0.0) + cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(cap_item_seqlens): + cap_attn_mask[i, :seq_len] = 1 + + # Run context refiner blocks + for layer in module.context_refiner: + cap_batched = layer(cap_batched, cap_attn_mask, cap_cos_batched, cap_sin_batched) + + # Create unified sequence (image + caption) + unified_list = [] + unified_cos_list = [] + unified_sin_list = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified_list.append(torch.cat([x_batched[i][:x_len], cap_batched[i][:cap_len]])) + unified_cos_list.append(torch.cat([x_cos_batched[i][:x_len], cap_cos_batched[i][:cap_len]])) + unified_sin_list.append(torch.cat([x_sin_batched[i][:x_len], cap_sin_batched[i][:cap_len]])) + + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + unified_max_item_seqlen = max(unified_item_seqlens) + + unified = pad_sequence(unified_list, batch_first=True, padding_value=0.0) + unified_cos = pad_sequence(unified_cos_list, batch_first=True, padding_value=0.0) + unified_sin = pad_sequence(unified_sin_list, batch_first=True, padding_value=0.0) + unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + # ============================================================================ + # EXTRACT MODULATED INPUT (for cache decision) + # ============================================================================ + # Use the first main transformer block's modulation + # The main layers have modulation=True and process the unified sequence + block = module.layers[0] + # Get modulation parameters: scale_msa, gate_msa, scale_mlp, gate_mlp + mod_params = block.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) + scale_msa = 1.0 + mod_params[0] + # Extract modulated input: normalized hidden states scaled by modulation + modulated_input = block.attention_norm1(unified) * scale_msa + + # ============================================================================ + # DEFINE TRANSFORMER EXECUTION (Z-Image specific) + # ============================================================================ + def run_transformer_blocks(): + """Execute all Z-Image main transformer blocks.""" + h = unified + for layer in module.layers: + h = layer(h, unified_attn_mask, unified_cos, unified_sin, adaln_input) + return (h,) + + # ============================================================================ + # DEFINE POSTPROCESSING (Z-Image specific) + # ============================================================================ + def postprocess(h): + """Apply Z-Image specific output postprocessing.""" + h = module.all_final_layer[f"{patch_size}-{f_patch_size}"](h, adaln_input) + h = list(h.unbind(dim=0)) + output = module.unpatchify(h, x_size, patch_size, f_patch_size) + return output, {} + + # ============================================================================ + # RETURN CONTEXT + # ============================================================================ + return CacheContext( + modulated_input=modulated_input, + hidden_states=unified, + encoder_hidden_states=None, # Z-Image uses unified sequence, no separate encoder states + temb=adaln_input, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + extra_states={ + "unified_attn_mask": unified_attn_mask, + "unified_cos": unified_cos, + "unified_sin": unified_sin, + "x_size": x_size, + "x_item_seqlens": x_item_seqlens, + "patch_size": patch_size, + "f_patch_size": f_patch_size, + }, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -262,6 +448,7 @@ def extract_qwen_context( # on the transformer module and multiple pipelines can share the same transformer. EXTRACTOR_REGISTRY: dict[str, Callable] = { "QwenImageTransformer2DModel": extract_qwen_context, + "ZImageTransformer2DModel": extract_zimage_context, # Future models: # "FluxTransformer2DModel": extract_flux_context, # "CogVideoXTransformer3DModel": extract_cogvideox_context, diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index cf04bf172..247d914e9 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -33,6 +33,7 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.base import CachedTransformer from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, @@ -520,7 +521,7 @@ class RopeEmbedder: return torch.cat(cos_result, dim=-1), torch.cat(sin_result, dim=-1) -class ZImageTransformer2DModel(nn.Module): +class ZImageTransformer2DModel(CachedTransformer): """Z-Image Transformer model for image generation. Sequence Parallelism: