d150a4fde7
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Cursor <cursoragent@cursor.com>
1490 lines
59 KiB
Python
1490 lines
59 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Model-specific extractors for TeaCache.
|
|
|
|
This module provides a registry of extractor functions that know how to extract
|
|
modulated inputs from different transformer architectures. Adding support for
|
|
a new model requires only adding a new extractor function to the registry.
|
|
|
|
With Option B enhancement, extractors now return a CacheContext object containing
|
|
all model-specific information needed for generic caching, including preprocessing,
|
|
transformer execution, and postprocessing logic.
|
|
"""
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.forward_context import get_forward_context
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class CacheContext:
|
|
"""
|
|
Context object containing all model-specific information for caching.
|
|
|
|
This allows the TeaCacheHook to remain completely generic - all model-specific
|
|
logic is encapsulated in the extractor that returns this context.
|
|
|
|
Attributes:
|
|
modulated_input: Tensor used for cache decision (similarity comparison).
|
|
Must be a torch.Tensor extracted from the first transformer block,
|
|
typically after applying normalization and modulation.
|
|
|
|
hidden_states: Current hidden states (will be modified by caching).
|
|
Must be a torch.Tensor representing the main image/latent states
|
|
after preprocessing but before transformer blocks.
|
|
|
|
encoder_hidden_states: Optional encoder states (for dual-stream models).
|
|
Set to None for single-stream models (e.g., Flux).
|
|
For dual-stream models (e.g., Qwen), contains text encoder outputs.
|
|
|
|
temb: Timestep embedding tensor.
|
|
Must be a torch.Tensor containing the timestep conditioning.
|
|
|
|
run_transformer_blocks: Callable that executes model-specific transformer blocks.
|
|
Signature: () -> tuple[torch.Tensor, ...]
|
|
|
|
Returns:
|
|
tuple containing:
|
|
- [0]: processed hidden_states (required)
|
|
- [1]: processed encoder_hidden_states (optional, only for dual-stream)
|
|
|
|
Example for single-stream:
|
|
def run_blocks():
|
|
h = hidden_states
|
|
for block in module.transformer_blocks:
|
|
h = block(h, temb=temb)
|
|
return (h,)
|
|
|
|
Example for dual-stream:
|
|
def run_blocks():
|
|
h, e = hidden_states, encoder_hidden_states
|
|
for block in module.transformer_blocks:
|
|
e, h = block(h, e, temb=temb)
|
|
return (h, e)
|
|
|
|
postprocess: Callable that does model-specific output postprocessing.
|
|
Signature: (torch.Tensor) -> Union[torch.Tensor, Transformer2DModelOutput, tuple]
|
|
|
|
Takes the processed hidden_states and applies final transformations
|
|
(normalization, projection) to produce the model output.
|
|
|
|
Example:
|
|
def postprocess(h):
|
|
h = module.norm_out(h, temb)
|
|
output = module.proj_out(h)
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
extra_states: Optional dict for additional model-specific state.
|
|
Use this for models that need to pass additional context beyond
|
|
the standard fields.
|
|
"""
|
|
|
|
modulated_input: torch.Tensor
|
|
hidden_states: torch.Tensor
|
|
encoder_hidden_states: torch.Tensor | None
|
|
temb: torch.Tensor
|
|
run_transformer_blocks: Callable[[], tuple[torch.Tensor, ...]]
|
|
postprocess: Callable[[torch.Tensor], Any]
|
|
extra_states: dict[str, Any] | None = None
|
|
|
|
def validate(self) -> None:
|
|
"""
|
|
Validate that the CacheContext contains valid data.
|
|
|
|
Raises:
|
|
TypeError: If fields have wrong types
|
|
ValueError: If tensors have invalid properties
|
|
RuntimeError: If callables fail basic invocation tests
|
|
|
|
This method should be called after creating a CacheContext to catch
|
|
common developer errors early with clear error messages.
|
|
"""
|
|
# Validate tensor fields
|
|
if not isinstance(self.modulated_input, torch.Tensor):
|
|
raise TypeError(f"modulated_input must be torch.Tensor, got {type(self.modulated_input)}")
|
|
|
|
if not isinstance(self.hidden_states, torch.Tensor):
|
|
raise TypeError(f"hidden_states must be torch.Tensor, got {type(self.hidden_states)}")
|
|
|
|
if self.encoder_hidden_states is not None and not isinstance(self.encoder_hidden_states, torch.Tensor):
|
|
raise TypeError(
|
|
f"encoder_hidden_states must be torch.Tensor or None, got {type(self.encoder_hidden_states)}"
|
|
)
|
|
|
|
if not isinstance(self.temb, torch.Tensor):
|
|
raise TypeError(f"temb must be torch.Tensor, got {type(self.temb)}")
|
|
|
|
# Validate callables
|
|
if not callable(self.run_transformer_blocks):
|
|
raise TypeError(f"run_transformer_blocks must be callable, got {type(self.run_transformer_blocks)}")
|
|
|
|
if not callable(self.postprocess):
|
|
raise TypeError(f"postprocess must be callable, got {type(self.postprocess)}")
|
|
|
|
# Validate tensor shapes are compatible
|
|
if self.modulated_input.shape[0] != self.hidden_states.shape[0]:
|
|
raise ValueError(
|
|
f"Batch size mismatch: modulated_input has batch size "
|
|
f"{self.modulated_input.shape[0]}, but hidden_states has "
|
|
f"{self.hidden_states.shape[0]}"
|
|
)
|
|
|
|
# Validate devices match
|
|
if self.modulated_input.device != self.hidden_states.device:
|
|
raise ValueError(
|
|
f"Device mismatch: modulated_input on {self.modulated_input.device}, "
|
|
f"hidden_states on {self.hidden_states.device}"
|
|
)
|
|
|
|
|
|
def extract_qwen_context(
|
|
module: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
encoder_hidden_states_mask: torch.Tensor,
|
|
timestep: torch.Tensor | float | int,
|
|
img_shapes: torch.Tensor,
|
|
txt_seq_lens: torch.Tensor,
|
|
guidance: torch.Tensor | None = None,
|
|
additional_t_cond: torch.Tensor | None = None,
|
|
attention_kwargs: dict[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for QwenImageTransformer2DModel.
|
|
|
|
This is the ONLY Qwen-specific code needed for TeaCache support.
|
|
It encapsulates preprocessing, modulated input extraction, transformer execution,
|
|
and postprocessing logic.
|
|
|
|
Args:
|
|
module: QwenImageTransformer2DModel instance
|
|
hidden_states: Input hidden states tensor
|
|
encoder_hidden_states: Text encoder outputs
|
|
encoder_hidden_states_mask: Mask for text encoder
|
|
timestep: Current diffusion timestep
|
|
img_shapes: Image shapes for position embedding
|
|
txt_seq_lens: Text sequence lengths
|
|
guidance: Optional guidance scale for CFG
|
|
additional_t_cond: Optional additional timestep conditioning
|
|
attention_kwargs: Additional attention arguments
|
|
**kwargs: Additional keyword arguments ignored by this extractor
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
|
raise ValueError("Module must have transformer_blocks")
|
|
|
|
# ============================================================================
|
|
# PREPROCESSING (Qwen-specific)
|
|
# ============================================================================
|
|
# Call image_rope_prepare instead of img_in + pos_embed directly.
|
|
# This ensures the SequenceParallelSplitHook registered on image_rope_prepare
|
|
# fires when SP is enabled, correctly sharding hidden_states and vid_freqs.
|
|
hidden_states, vid_freqs, txt_freqs = module.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens)
|
|
image_rotary_emb = (vid_freqs, txt_freqs)
|
|
|
|
timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
|
|
|
|
# Call modulate_index_prepare instead of handling timestep directly.
|
|
# For zero_cond_t=False: timestep unchanged, modulate_index=None.
|
|
# For zero_cond_t=True: timestep is doubled, modulate_index is created and
|
|
# sharded by the SequenceParallelSplitHook on modulate_index_prepare so that
|
|
# its sequence dimension matches the already-sharded hidden_states.
|
|
timestep, modulate_index = module.modulate_index_prepare(timestep, img_shapes)
|
|
|
|
encoder_hidden_states = module.txt_norm(encoder_hidden_states)
|
|
encoder_hidden_states = module.txt_in(encoder_hidden_states)
|
|
|
|
if guidance is not None:
|
|
guidance = guidance.to(hidden_states.dtype) * 1000
|
|
|
|
temb = (
|
|
module.time_text_embed(timestep, hidden_states, additional_t_cond)
|
|
if guidance is None
|
|
else module.time_text_embed(timestep, guidance, hidden_states, additional_t_cond)
|
|
)
|
|
|
|
# ============================================================================
|
|
# EXTRACT MODULATED INPUT (for cache decision)
|
|
# ============================================================================
|
|
block = module.transformer_blocks[0]
|
|
img_mod_params = block.img_mod(temb)
|
|
img_mod1, _ = img_mod_params.chunk(2, dim=-1)
|
|
img_scale1, img_shift1, _ = block._modulate(img_mod1)
|
|
img_modulated = block.img_norm1(hidden_states, img_scale1, img_shift1)
|
|
|
|
# ============================================================================
|
|
# DEFINE TRANSFORMER EXECUTION (Qwen-specific)
|
|
# ============================================================================
|
|
def run_transformer_blocks():
|
|
"""Execute all Qwen transformer blocks."""
|
|
h = hidden_states
|
|
e = encoder_hidden_states
|
|
encoder_mask = encoder_hidden_states_mask
|
|
hidden_states_mask = None # default
|
|
if module.parallel_config is not None and module.parallel_config.sequence_parallel_size > 1:
|
|
ctx = get_forward_context()
|
|
if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
|
|
# Create mask for the full (padded) sequence
|
|
# valid positions = True, padding positions = False
|
|
batch_size = hidden_states.shape[0]
|
|
padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
|
|
hidden_states_mask = torch.ones(
|
|
batch_size,
|
|
padded_seq_len,
|
|
dtype=torch.bool,
|
|
device=hidden_states.device,
|
|
)
|
|
hidden_states_mask[:, ctx.sp_original_seq_len :] = False
|
|
|
|
# if mask is all true, set it to None
|
|
if hidden_states_mask is not None and hidden_states_mask.all():
|
|
hidden_states_mask = None
|
|
if encoder_mask is not None and encoder_mask.all():
|
|
encoder_mask = None
|
|
for block in module.transformer_blocks:
|
|
e, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
encoder_hidden_states_mask=encoder_mask,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
joint_attention_kwargs=attention_kwargs,
|
|
modulate_index=modulate_index,
|
|
hidden_states_mask=hidden_states_mask,
|
|
)
|
|
return (h, e)
|
|
|
|
# ============================================================================
|
|
# DEFINE POSTPROCESSING (Qwen-specific)
|
|
# ============================================================================
|
|
return_dict = kwargs.get("return_dict", True)
|
|
|
|
def postprocess(h):
|
|
"""Apply Qwen-specific output postprocessing."""
|
|
if getattr(module, "zero_cond_t", False):
|
|
h = module.norm_out(h, temb.chunk(2, dim=0)[0])
|
|
else:
|
|
h = module.norm_out(h, temb)
|
|
output = module.proj_out(h)
|
|
if not return_dict:
|
|
return (output,)
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
# ============================================================================
|
|
# RETURN CONTEXT
|
|
# ============================================================================
|
|
return CacheContext(
|
|
modulated_input=img_modulated,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_bagel_context(
|
|
module: nn.Module,
|
|
x_t: torch.Tensor,
|
|
timestep: torch.Tensor | float | int,
|
|
packed_vae_token_indexes: torch.LongTensor,
|
|
packed_vae_position_ids: torch.LongTensor,
|
|
packed_text_ids: torch.LongTensor,
|
|
packed_text_indexes: torch.LongTensor,
|
|
packed_position_ids: torch.LongTensor,
|
|
packed_seqlens: torch.IntTensor,
|
|
past_key_values: Any,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for Bagel model.
|
|
|
|
Args:
|
|
module: Bagel instance
|
|
x_t: Latent image input
|
|
timestep: Current timestep
|
|
packed_vae_token_indexes: Indexes for VAE tokens in packed sequence
|
|
packed_vae_position_ids: Position IDs for VAE tokens
|
|
packed_text_ids: Text token IDs
|
|
packed_text_indexes: Indexes for text tokens in packed sequence
|
|
packed_position_ids: Global position IDs
|
|
packed_seqlens: Sequence lengths
|
|
past_key_values: KV cache
|
|
**kwargs: Additional keyword arguments
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
|
|
# 1. Embed text
|
|
packed_text_embedding = module.language_model.model.embed_tokens(packed_text_ids)
|
|
packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), module.hidden_size))
|
|
packed_sequence[packed_text_indexes] = packed_text_embedding
|
|
|
|
# 2. Embed timestep
|
|
if not isinstance(timestep, torch.Tensor):
|
|
timestep = torch.tensor([timestep], device=x_t.device)
|
|
if timestep.dim() == 0:
|
|
timestep = timestep.unsqueeze(0)
|
|
|
|
# 3. Embed image (x_t)
|
|
packed_pos_embed = module.latent_pos_embed(packed_vae_position_ids)
|
|
packed_timestep_embeds = module.time_embedder(timestep)
|
|
|
|
x_t_emb = module.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
|
|
if x_t_emb.dtype != packed_sequence.dtype:
|
|
x_t_emb = x_t_emb.to(packed_sequence.dtype)
|
|
|
|
packed_sequence[packed_vae_token_indexes] = x_t_emb
|
|
|
|
# Use the full packed sequence as modulated input to match hidden_states size
|
|
modulated_input = packed_sequence
|
|
|
|
def run_transformer_blocks():
|
|
extra_inputs = {}
|
|
if module.use_moe:
|
|
extra_inputs = {
|
|
"mode": "gen",
|
|
"packed_vae_token_indexes": packed_vae_token_indexes,
|
|
"packed_text_indexes": packed_text_indexes,
|
|
}
|
|
|
|
output = module.language_model.forward(
|
|
packed_query_sequence=packed_sequence,
|
|
query_lens=packed_seqlens,
|
|
packed_query_position_ids=packed_position_ids,
|
|
past_key_values=past_key_values,
|
|
update_past_key_values=False,
|
|
is_causal=False,
|
|
**extra_inputs,
|
|
)
|
|
return (output.packed_query_sequence,)
|
|
|
|
def postprocess(h):
|
|
v_t = module.llm2vae(h)
|
|
v_t = v_t[packed_vae_token_indexes]
|
|
return v_t
|
|
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=packed_sequence, # Use full packed sequence
|
|
encoder_hidden_states=None,
|
|
temb=packed_timestep_embeds, # Approximate
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
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,
|
|
},
|
|
)
|
|
|
|
|
|
def extract_flux2_klein_context(
|
|
module: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
timestep: torch.LongTensor = None,
|
|
img_ids: torch.Tensor = None,
|
|
txt_ids: torch.Tensor = None,
|
|
guidance: torch.Tensor | None = None,
|
|
joint_attention_kwargs: dict[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for Flux2Klein model.
|
|
|
|
Caches the full transformer output (including single_transformer_blocks).
|
|
When cache is reused, single_transformer_blocks is skipped to achieve maximum speedup.
|
|
|
|
Args:
|
|
module: Flux2Transformer2DModel instance
|
|
hidden_states: Input image hidden states tensor
|
|
encoder_hidden_states: Input text hidden states tensor
|
|
timestep: Current diffusion timestep
|
|
img_ids: Image position IDs for RoPE
|
|
txt_ids: Text position IDs for RoPE
|
|
guidance: Optional guidance scale for CFG
|
|
joint_attention_kwargs: Additional attention kwargs
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
|
raise ValueError("Module must have transformer_blocks")
|
|
|
|
# ============================================================================
|
|
# PREPROCESSING (Flux2-specific)
|
|
# ============================================================================
|
|
dtype = hidden_states.dtype
|
|
|
|
num_txt_tokens = encoder_hidden_states.shape[1]
|
|
|
|
timestep = timestep.to(dtype=dtype) * 1000
|
|
if guidance is not None:
|
|
guidance = guidance.to(dtype=dtype) * 1000
|
|
|
|
temb = module.time_guidance_embed(timestep, guidance)
|
|
|
|
double_stream_mod_img = module.double_stream_modulation_img(temb)
|
|
double_stream_mod_txt = module.double_stream_modulation_txt(temb)
|
|
single_stream_mod = module.single_stream_modulation(temb)[0]
|
|
|
|
hidden_states = module.x_embedder(hidden_states)
|
|
encoder_hidden_states = module.context_embedder(encoder_hidden_states)
|
|
|
|
if img_ids.ndim == 3:
|
|
img_ids = img_ids[0]
|
|
if txt_ids.ndim == 3:
|
|
txt_ids = txt_ids[0]
|
|
|
|
image_rotary_emb = module.pos_embed(img_ids)
|
|
text_rotary_emb = module.pos_embed(txt_ids)
|
|
concat_rotary_emb = (
|
|
torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
|
|
torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
|
|
)
|
|
|
|
# ============================================================================
|
|
# EXTRACT MODULATED INPUT (for cache decision)
|
|
# ============================================================================
|
|
block = module.transformer_blocks[0]
|
|
|
|
norm_hidden_states = block.norm1(hidden_states)
|
|
norm_hidden_states = (1 + double_stream_mod_img[0][1]) * norm_hidden_states + double_stream_mod_img[0][0]
|
|
|
|
modulated_input = norm_hidden_states
|
|
|
|
# ============================================================================
|
|
# DEFINE TRANSFORMER EXECUTION (Flux2-specific)
|
|
# ============================================================================
|
|
def run_flux2_transformer_blocks():
|
|
h = hidden_states
|
|
c = encoder_hidden_states
|
|
for block in module.transformer_blocks:
|
|
c, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=c,
|
|
temb_mod_params_img=double_stream_mod_img,
|
|
temb_mod_params_txt=double_stream_mod_txt,
|
|
image_rotary_emb=concat_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
return (h, c)
|
|
|
|
def run_flux2_full_transformer_with_single(ori_h, ori_c):
|
|
h = ori_h
|
|
c = ori_c
|
|
for block in module.transformer_blocks:
|
|
c, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=c,
|
|
temb_mod_params_img=double_stream_mod_img,
|
|
temb_mod_params_txt=double_stream_mod_txt,
|
|
image_rotary_emb=concat_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
h_concat = torch.cat([c, h], dim=1)
|
|
for block in module.single_transformer_blocks:
|
|
h_concat = block(
|
|
hidden_states=h_concat,
|
|
encoder_hidden_states=None,
|
|
temb_mod_params=single_stream_mod,
|
|
image_rotary_emb=concat_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
final_hidden_states = h_concat[:, num_txt_tokens:, ...]
|
|
return final_hidden_states, c
|
|
|
|
# ============================================================================
|
|
# DEFINE POSTPROCESSING (Flux2-specific)
|
|
# ============================================================================
|
|
return_dict = kwargs.get("return_dict", True)
|
|
|
|
def postprocess(h):
|
|
h = module.norm_out(h, temb)
|
|
h = module.proj_out(h)
|
|
if not return_dict:
|
|
return (h,)
|
|
return Transformer2DModelOutput(sample=h)
|
|
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
run_transformer_blocks=run_flux2_transformer_blocks,
|
|
postprocess=postprocess,
|
|
extra_states={
|
|
"run_flux2_full_transformer_with_single": run_flux2_full_transformer_with_single,
|
|
},
|
|
)
|
|
|
|
|
|
def extract_longcat_context(
|
|
module: nn.Module, # LongCatImageTransformer2DModel
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
timestep=None,
|
|
img_ids=None,
|
|
txt_ids=None,
|
|
guidance=None,
|
|
**kwargs,
|
|
) -> CacheContext:
|
|
"""Extract the cache context for LongCat Image.
|
|
|
|
Similar to other extractors, this is currently the only code needed
|
|
for TeaCache support for LongCat image, and encapsulates preprocessing,
|
|
modulated input extraction, transformer execution, and postprocessing
|
|
logic.
|
|
|
|
Args & kawrgs are identical to the inputs to LongCat Image's forward.
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
# TODO (Alex) - Refactor TeaCache extractors to more tightly integrate with .forward
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
# 1. Model specific preprocessing
|
|
fwd_context = get_forward_context()
|
|
sp_size = module.parallel_config.sequence_parallel_size
|
|
if sp_size is not None and sp_size > 1:
|
|
# NOTE: For now, we set this to False on the forward context
|
|
# to be consistent with LongCat Image's current behavior when
|
|
# TeaCache is enabled. We do not need to reset it in post process
|
|
# since we should never split text embed in sp for this model.
|
|
fwd_context.split_text_embed_in_sp = False
|
|
|
|
hidden_states = module.x_embedder(hidden_states)
|
|
|
|
timestep = timestep.to(hidden_states.dtype) * 1000
|
|
|
|
temb = module.time_embed(timestep, hidden_states.dtype)
|
|
encoder_hidden_states = module.context_embedder(encoder_hidden_states)
|
|
|
|
# Compute RoPE embeddings via rope_preparer module
|
|
# _sp_plan will automatically shard img_cos/img_sin (outputs 2, 3)
|
|
# txt_cos/txt_sin (outputs 0, 1) remain replicated for dual-stream attention
|
|
txt_cos, txt_sin, img_cos, img_sin = module.rope_preparer(txt_ids, img_ids)
|
|
|
|
# Reconstruct image_rotary_emb with chunked values
|
|
# Final shape: (txt_seq_len + img_seq_len // SP, head_dim)
|
|
image_rotary_emb = (
|
|
torch.cat([txt_cos, img_cos], dim=0),
|
|
torch.cat([txt_sin, img_sin], dim=0),
|
|
)
|
|
|
|
# 2. Extract the modulated output from the first mm-DiT block
|
|
first_block = module.transformer_blocks[0]
|
|
img_modulated = first_block.norm1(hidden_states, emb=temb)[0]
|
|
|
|
# 3. Define the transformer execution
|
|
def run_transformer_blocks():
|
|
"""Execute all Longcat transformer blocks."""
|
|
h = hidden_states
|
|
e = encoder_hidden_states
|
|
for block in module.transformer_blocks:
|
|
e, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
)
|
|
|
|
for block in module.single_transformer_blocks:
|
|
e, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
)
|
|
# Hook expects hidden states to be first
|
|
return (h, e)
|
|
|
|
# 4. Postprocessing
|
|
def postprocess(h):
|
|
"""Apply Longcat-specific output postprocessing."""
|
|
h = module.norm_out(h, temb)
|
|
output = module.proj_out(h)
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
# 5. Return the CacheContext
|
|
return CacheContext(
|
|
modulated_input=img_modulated,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_stable_audio_context(
|
|
module: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
global_hidden_states: torch.Tensor | None = None,
|
|
rotary_embedding: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
return_dict: bool = True,
|
|
attention_mask: torch.Tensor | None = None,
|
|
encoder_attention_mask: torch.Tensor | None = None,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for Stable Audio DiT model.
|
|
|
|
Architecture Notes:
|
|
- Stable Audio uses standard LayerNorm
|
|
- Timestep conditioning via global_hidden_states prepended to sequence
|
|
- Single-stream model (cross-attention handled separately)
|
|
- Input: [B, C, L] (C=in_channels) -> transpose -> [B, L, C] -> project -> [B, L, inner_dim]
|
|
- Global states prepended: [B, 1+L, inner_dim]
|
|
"""
|
|
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
|
raise ValueError("Module must have transformer_blocks attribute with at least one block")
|
|
|
|
# Cast inputs to match model weights dtype
|
|
hidden_states = hidden_states.to(module.dtype)
|
|
encoder_hidden_states = encoder_hidden_states.to(module.dtype)
|
|
if global_hidden_states is not None:
|
|
global_hidden_states = global_hidden_states.to(module.dtype)
|
|
|
|
cross_attention_hidden_states = module.cross_attention_proj(encoder_hidden_states)
|
|
global_hidden_states = module.global_proj(global_hidden_states)
|
|
time_hidden_states = module.timestep_proj(module.time_proj(timestep.to(module.dtype)))
|
|
|
|
# Combine global and time embeddings → [B, 1, inner_dim]
|
|
temb = global_hidden_states + time_hidden_states.unsqueeze(1)
|
|
|
|
hidden_states = module.preprocess_conv(hidden_states) + hidden_states
|
|
hidden_states = hidden_states.transpose(1, 2)
|
|
|
|
# Capture the original sequence length for safe post-processing slicing
|
|
original_seq_len = hidden_states.shape[1]
|
|
|
|
hidden_states = module.proj_in(hidden_states)
|
|
hidden_states = torch.cat([temb, hidden_states], dim=1)
|
|
|
|
# attention_mask is not used in the standard Stable Audio Open inference path.
|
|
assert attention_mask is None, (
|
|
"attention_mask is not supported in extract_stable_audio_context; expected None for Stable Audio Open."
|
|
)
|
|
|
|
# Stable Audio prepends the combined global+time embedding (`temb`) to the sequence.
|
|
# Therefore, the standard LayerNorm applied here still captures the timestep signal
|
|
# within the first token of the output, giving the cache discriminator the info it needs.
|
|
first_block = module.transformer_blocks[0]
|
|
modulated_input = first_block.norm1(hidden_states)
|
|
|
|
def run_transformer_blocks() -> tuple[torch.Tensor]:
|
|
"""
|
|
Execute all Stable Audio transformer blocks.
|
|
|
|
Returns:
|
|
Tuple containing only hidden_states (single-stream model).
|
|
Format: (hidden_states,)
|
|
"""
|
|
h = hidden_states
|
|
for block in module.transformer_blocks:
|
|
h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=cross_attention_hidden_states,
|
|
rotary_embedding=rotary_embedding,
|
|
attention_mask=attention_mask,
|
|
encoder_attention_mask=encoder_attention_mask,
|
|
)
|
|
return (h,)
|
|
|
|
def postprocess(h: torch.Tensor) -> Any:
|
|
"""
|
|
Apply Stable Audio-specific output postprocessing.
|
|
|
|
Args:
|
|
h: Hidden states from transformer blocks [B, 1+L, inner_dim]
|
|
|
|
Returns:
|
|
Transformer2DModelOutput or tuple based on return_dict
|
|
"""
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
h = module.proj_out(h)
|
|
h = h.transpose(1, 2)[:, :, -original_seq_len:]
|
|
output = module.postprocess_conv(h) + h
|
|
if return_dict:
|
|
return Transformer2DModelOutput(sample=output)
|
|
return (output,)
|
|
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=None,
|
|
temb=temb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_flux2_context(
|
|
module: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor = None,
|
|
timestep: torch.LongTensor = None,
|
|
img_ids: torch.Tensor = None,
|
|
txt_ids: torch.Tensor = None,
|
|
guidance: torch.Tensor | None = None,
|
|
joint_attention_kwargs: dict[str, Any] | None = None,
|
|
return_dict: bool = True,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for Flux2Transformer2DModel.
|
|
|
|
This is the ONLY Flux2-specific code needed for TeaCache support.
|
|
It encapsulates preprocessing, modulated input extraction, transformer execution,
|
|
and postprocessing logic.
|
|
|
|
Args:
|
|
module: Flux2Transformer2DModel instance
|
|
hidden_states: Input hidden states tensor
|
|
encoder_hidden_states: Text encoder outputs
|
|
timestep: Current diffusion timestep
|
|
img_ids: Image inputs for position embedding
|
|
txt_ids: Text inputs for position embedding
|
|
guidance: Optional guidance scale for CFG
|
|
joint_attention_kwargs: Additional attention arguments
|
|
return_dict: Whether to return a Transformer2DModelOutput instead of a plain tensor
|
|
**kwargs: Additional keyword arguments ignored by this extractor
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
|
raise ValueError("Module must have transformer_blocks")
|
|
|
|
# ============================================================================
|
|
# PREPROCESSING (Flux2-specific)
|
|
# ============================================================================
|
|
num_txt_tokens = encoder_hidden_states.shape[1]
|
|
|
|
timestep = timestep.to(hidden_states.dtype) * 1000
|
|
if guidance is not None:
|
|
guidance = guidance.to(hidden_states.dtype) * 1000
|
|
|
|
temb = module.time_guidance_embed(timestep, guidance)
|
|
|
|
double_stream_mod_img = module.double_stream_modulation_img(temb)
|
|
double_stream_mod_txt = module.double_stream_modulation_txt(temb)
|
|
single_stream_mod = module.single_stream_modulation(temb)[0]
|
|
|
|
hidden_states = module.x_embedder(hidden_states)
|
|
encoder_hidden_states = module.context_embedder(encoder_hidden_states)
|
|
|
|
if img_ids.ndim == 3:
|
|
img_ids = img_ids[0]
|
|
if txt_ids.ndim == 3:
|
|
txt_ids = txt_ids[0]
|
|
|
|
if current_omni_platform.is_npu():
|
|
freqs_cos_image, freqs_sin_image = module.pos_embed(img_ids.cpu())
|
|
image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu())
|
|
freqs_cos_text, freqs_sin_text = module.pos_embed(txt_ids.cpu())
|
|
text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu())
|
|
else:
|
|
image_rotary_emb = module.pos_embed(img_ids)
|
|
text_rotary_emb = module.pos_embed(txt_ids)
|
|
concat_rotary_emb = (
|
|
torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
|
|
torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
|
|
)
|
|
|
|
# ============================================================================
|
|
# EXTRACT MODULATED INPUT (for cache decision)
|
|
# ============================================================================
|
|
block = module.transformer_blocks[0]
|
|
(shift_msa, scale_msa, gate_msa), _ = double_stream_mod_img
|
|
modulated_input = block.norm1(hidden_states)
|
|
modulated_input = (1 + scale_msa) * modulated_input + shift_msa
|
|
|
|
# ============================================================================
|
|
# DEFINE TRANSFORMER EXECUTION (Flux2-specific)
|
|
# ============================================================================
|
|
def run_transformer_blocks():
|
|
"""Execute all Flux2 transformer blocks."""
|
|
h = hidden_states
|
|
e = encoder_hidden_states
|
|
|
|
for transformer_block in module.transformer_blocks:
|
|
e, h = transformer_block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
temb_mod_params_img=double_stream_mod_img,
|
|
temb_mod_params_txt=double_stream_mod_txt,
|
|
image_rotary_emb=concat_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
h = torch.cat([e, h], dim=1)
|
|
|
|
for single_transformer_block in module.single_transformer_blocks:
|
|
h = single_transformer_block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=None,
|
|
temb_mod_params=single_stream_mod,
|
|
image_rotary_emb=concat_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
|
|
h = h[:, num_txt_tokens:, ...]
|
|
return (h,)
|
|
|
|
# ============================================================================
|
|
# DEFINE POSTPROCESSING
|
|
# ============================================================================
|
|
def postprocess(h):
|
|
h = module.norm_out(h, temb)
|
|
output = module.proj_out(h)
|
|
if not return_dict:
|
|
return (output,)
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
# ============================================================================
|
|
# RETURN CONTEXT
|
|
# ============================================================================
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_flux_context(
|
|
module: nn.Module,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
pooled_projections: torch.Tensor,
|
|
timestep: torch.LongTensor,
|
|
img_ids: torch.Tensor,
|
|
txt_ids: torch.Tensor,
|
|
guidance: torch.Tensor | None = None,
|
|
joint_attention_kwargs: dict[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""
|
|
Extract cache context for FluxTransformer2DModel.
|
|
|
|
This mirrors the standard FLUX.1 transformer forward path while exposing
|
|
the first modulated image stream tensor as TeaCache's similarity signal.
|
|
|
|
Args:
|
|
module: FluxTransformer2DModel instance
|
|
hidden_states: Input image hidden states tensor
|
|
encoder_hidden_states: Text encoder outputs
|
|
pooled_projections: Pooled text conditioning
|
|
timestep: Current diffusion timestep
|
|
img_ids: Image position IDs for RoPE
|
|
txt_ids: Text position IDs for RoPE
|
|
guidance: Optional guidance scale for guidance-distilled variants
|
|
joint_attention_kwargs: Additional attention kwargs
|
|
**kwargs: Additional keyword arguments
|
|
|
|
Returns:
|
|
CacheContext with all information needed for generic caching
|
|
"""
|
|
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
|
|
|
if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
|
|
raise ValueError("Module must have transformer_blocks")
|
|
|
|
# ============================================================================
|
|
# PREPROCESSING (Flux1-specific)
|
|
# ============================================================================
|
|
hidden_states = module.x_embedder(hidden_states)
|
|
timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
|
|
|
|
if guidance is not None:
|
|
guidance = guidance.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
|
|
|
|
temb = (
|
|
module.time_text_embed(timestep, pooled_projections)
|
|
if guidance is None
|
|
else module.time_text_embed(timestep, guidance, pooled_projections)
|
|
)
|
|
encoder_hidden_states = module.context_embedder(encoder_hidden_states)
|
|
|
|
if txt_ids.ndim == 3:
|
|
txt_ids = txt_ids[0]
|
|
if img_ids.ndim == 3:
|
|
img_ids = img_ids[0]
|
|
|
|
ids = torch.cat((txt_ids, img_ids), dim=0)
|
|
if current_omni_platform.is_npu():
|
|
freqs_cos, freqs_sin = module.pos_embed(ids.cpu())
|
|
image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
|
|
else:
|
|
image_rotary_emb = module.pos_embed(ids)
|
|
|
|
# ============================================================================
|
|
# EXTRACT MODULATED INPUT (for cache decision)
|
|
# ============================================================================
|
|
first_block = module.transformer_blocks[0]
|
|
modulated_input, *_ = first_block.norm1(hidden_states, emb=temb)
|
|
|
|
# ============================================================================
|
|
# DEFINE TRANSFORMER EXECUTION (Flux1-specific)
|
|
# ============================================================================
|
|
def run_transformer_blocks() -> tuple[torch.Tensor, torch.Tensor]:
|
|
h = hidden_states
|
|
e = encoder_hidden_states
|
|
|
|
for block in module.transformer_blocks:
|
|
e, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
|
|
for block in module.single_transformer_blocks:
|
|
e, h = block(
|
|
hidden_states=h,
|
|
encoder_hidden_states=e,
|
|
temb=temb,
|
|
image_rotary_emb=image_rotary_emb,
|
|
joint_attention_kwargs=joint_attention_kwargs,
|
|
)
|
|
|
|
return (h, e)
|
|
|
|
return_dict = kwargs.get("return_dict", True)
|
|
|
|
# ============================================================================
|
|
# DEFINE POSTPROCESSING
|
|
# ============================================================================
|
|
def postprocess(h: torch.Tensor) -> Any:
|
|
h = module.norm_out(h, temb)
|
|
output = module.proj_out(h)
|
|
if not return_dict:
|
|
return (output,)
|
|
return Transformer2DModelOutput(sample=output)
|
|
|
|
# ============================================================================
|
|
# RETURN CONTEXT
|
|
# ============================================================================
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
temb=temb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_sensenova_u1_context(
|
|
module: nn.Module,
|
|
input_ids: torch.Tensor | None = None,
|
|
indexes: torch.Tensor | None = None,
|
|
attention_mask: dict[str, torch.Tensor] | torch.Tensor | None = None,
|
|
past_key_values: Any | None = None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
use_cache: bool | None = None,
|
|
embed_only: bool = False,
|
|
compute_logits: bool = True,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""Extract cache context for SenseNovaU1ForCausalLM denoising forwards."""
|
|
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
|
|
SenseNovaU1CausalLMOutput,
|
|
)
|
|
|
|
layer_kwargs = dict(kwargs)
|
|
layer_kwargs.pop("cache_dit_skip", None)
|
|
image_gen_indicators = layer_kwargs.pop("image_gen_indicators", None)
|
|
exist_und = (~image_gen_indicators).any().item()
|
|
exist_gen = image_gen_indicators.any().item()
|
|
causal_mask_mapping = attention_mask
|
|
|
|
first_layer = module.model.layers[0]
|
|
modulated_input = first_layer.input_layernorm_mot_gen(inputs_embeds)
|
|
|
|
def run_transformer_blocks():
|
|
h = inputs_embeds
|
|
for layer in module.model.layers:
|
|
h = layer(
|
|
h,
|
|
image_gen_indicators=image_gen_indicators,
|
|
exist_und=exist_und,
|
|
exist_gen=exist_gen,
|
|
indexes=indexes,
|
|
attention_mask=causal_mask_mapping,
|
|
past_key_values=past_key_values,
|
|
**layer_kwargs,
|
|
)
|
|
return (h,)
|
|
|
|
def postprocess(h: torch.Tensor) -> SenseNovaU1CausalLMOutput:
|
|
h = module.model.norm_mot_gen(h)
|
|
|
|
logits = module.logits_processor(module.lm_head, h) if compute_logits else None
|
|
return SenseNovaU1CausalLMOutput(
|
|
logits=logits,
|
|
past_key_values=past_key_values if use_cache else None,
|
|
hidden_states=h,
|
|
)
|
|
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=inputs_embeds,
|
|
encoder_hidden_states=None,
|
|
temb=modulated_input,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
def extract_minimax_h3_context(
|
|
module: nn.Module,
|
|
**kwargs: Any,
|
|
) -> CacheContext:
|
|
"""Extract cache context for MiniMaxH3DiTModel.
|
|
|
|
Mirrors the packed multimodal forward in ``MiniMaxH3DiTModel.forward``:
|
|
preprocessing via ``_embed``, cacheable ``blocks`` loop, and
|
|
``final_layer`` postprocessing with row selection and update masks.
|
|
"""
|
|
from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
_FORWARD_SUPPORTED_KWARGS,
|
|
MINIMAX_H3_ADALN_MODALITY_NUM,
|
|
_build_rope_table,
|
|
_required_kwarg,
|
|
)
|
|
|
|
if not hasattr(module, "blocks") or len(module.blocks) == 0:
|
|
raise ValueError("Module must have blocks")
|
|
|
|
unexpected = sorted(set(kwargs) - _FORWARD_SUPPORTED_KWARGS)
|
|
if unexpected:
|
|
raise TypeError(
|
|
"MiniMaxH3DiTModel.forward received unexpected kwargs: "
|
|
f"{unexpected}; supported kwargs: "
|
|
f"{sorted(_FORWARD_SUPPORTED_KWARGS)}"
|
|
)
|
|
|
|
x = _required_kwarg(kwargs, "x")
|
|
audio_x = _required_kwarg(kwargs, "audio_x")
|
|
img_position_ids = _required_kwarg(kwargs, "img_position_ids")
|
|
unique_timesteps = _required_kwarg(kwargs, "unique_timesteps")
|
|
inverse_indices = _required_kwarg(kwargs, "inverse_indices").view(-1).to(torch.long)
|
|
update_mask = _required_kwarg(kwargs, "update_mask")
|
|
token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long)
|
|
skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False))
|
|
update_audio_mask = kwargs.get("update_audio_mask")
|
|
|
|
text_selected = _required_kwarg(kwargs, "prompt_embeds")
|
|
|
|
img_pos = module._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info")
|
|
audio_pos = module._pos_ids(_required_kwarg(kwargs, "audio_pos_info"), "audio_pos_info")
|
|
text_pos = module._pos_ids(
|
|
_required_kwarg(kwargs, "text_pos_info"),
|
|
"text_pos_info",
|
|
)
|
|
infer_out_pos = module._pos_ids(
|
|
_required_kwarg(kwargs, "img_pos_for_infer_output_info"),
|
|
"img_pos_for_infer_output_info",
|
|
)
|
|
|
|
psp = _required_kwarg(kwargs, "packed_seq_params")
|
|
cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
|
max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
|
|
# Mirror forward()'s host-side scalar so the refiner sees the number of
|
|
# packed requests without reading cu_seqlens off-device.
|
|
num_requests = int(module._psp_optional(psp, "num_requests", 1))
|
|
refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
|
|
refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
|
refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
|
|
video_layout = kwargs.get("video_token_layout")
|
|
|
|
if x.dim() != 3 or x.shape[0] != 1:
|
|
raise ValueError(f"x must be [1, S, C], got {list(x.shape)}")
|
|
seq_len = int(x.shape[1])
|
|
if token_tags.shape[0] != seq_len:
|
|
raise ValueError(f"token_tags must cover the full packed sequence ({seq_len}), got {token_tags.shape[0]}.")
|
|
if inverse_indices.shape[0] != seq_len:
|
|
raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
|
|
device = x.device
|
|
|
|
rope_table = _build_rope_table(module.rope(img_position_ids).to(device))
|
|
|
|
decoder_input, t_emb = module._embed(
|
|
x=x,
|
|
audio_x=audio_x,
|
|
text_embeddings_selected=text_selected,
|
|
unique_timesteps=unique_timesteps.view(-1).to(device),
|
|
img_pos=img_pos.to(device),
|
|
audio_pos=audio_pos.to(device),
|
|
text_pos=text_pos.to(device),
|
|
refiner_cu_seqlens=refiner_cu.to(device),
|
|
refiner_max_seqlen=refiner_max,
|
|
num_requests=num_requests,
|
|
seq_len=seq_len,
|
|
device=device,
|
|
local_span=(0, seq_len),
|
|
)
|
|
|
|
combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)
|
|
inverse_indices = inverse_indices.to(device)
|
|
cu_seqlens = cu_seqlens.to(device)
|
|
|
|
shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb)
|
|
modulated_input = module.blocks[0].norm1(decoder_input)
|
|
modulated_input = indexed_scale_shift_(
|
|
modulated_input,
|
|
shift_msa,
|
|
scale_msa,
|
|
combined_indices,
|
|
)
|
|
|
|
def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
|
|
hidden, block_rope, block_combined = module.sp_prepare(
|
|
decoder_input,
|
|
rope_table,
|
|
combined_indices,
|
|
)
|
|
for block in module.blocks:
|
|
hidden = block(
|
|
hidden,
|
|
t_emb=t_emb,
|
|
combined_indices=block_combined,
|
|
rope_table=block_rope,
|
|
cu_seqlens=cu_seqlens,
|
|
max_seqlen=max_seqlen,
|
|
packed_total=seq_len,
|
|
video_layout=video_layout,
|
|
)
|
|
hidden = module.sp_gather(hidden)
|
|
return (hidden,)
|
|
|
|
def postprocess(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
video_logits, audio_logits = module.final_layer(
|
|
hidden,
|
|
t_emb=t_emb,
|
|
inverse_indices=inverse_indices,
|
|
)
|
|
video_logits = video_logits.index_select(0, infer_out_pos.to(device))
|
|
audio_logits = audio_logits.index_select(0, audio_pos.to(device))
|
|
if not skip_mask_out_condition:
|
|
mask = update_mask.view(-1).to(device)
|
|
if mask.shape[0] != video_logits.shape[0]:
|
|
raise ValueError(f"update_mask length mismatch: {mask.shape[0]} != {video_logits.shape[0]}")
|
|
video_logits = video_logits * mask.unsqueeze(-1)
|
|
if update_audio_mask is not None:
|
|
audio_logits = audio_logits * update_audio_mask.view(-1).unsqueeze(-1)
|
|
return video_logits, audio_logits
|
|
|
|
return CacheContext(
|
|
modulated_input=modulated_input,
|
|
hidden_states=decoder_input,
|
|
encoder_hidden_states=None,
|
|
temb=t_emb,
|
|
run_transformer_blocks=run_transformer_blocks,
|
|
postprocess=postprocess,
|
|
)
|
|
|
|
|
|
# Registry for model-specific extractors
|
|
# Key: Transformer class name
|
|
# Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
|
|
#
|
|
# Note: Use the transformer class name as specified in pipelines as TeaCache hooks operate
|
|
# on the transformer module and multiple pipelines can share the same transformer.
|
|
EXTRACTOR_REGISTRY: dict[str, Callable] = {
|
|
"QwenImageTransformer2DModel": extract_qwen_context,
|
|
"Bagel": extract_bagel_context,
|
|
"ZImageTransformer2DModel": extract_zimage_context,
|
|
"Flux2Klein": extract_flux2_klein_context,
|
|
"StableAudioDiTModel": extract_stable_audio_context,
|
|
"Flux2Transformer2DModel": extract_flux2_context,
|
|
"LongCatImageTransformer2DModel": extract_longcat_context,
|
|
"FluxTransformer2DModel": extract_flux_context,
|
|
"SenseNovaU1ForCausalLM": extract_sensenova_u1_context,
|
|
"MiniMaxH3DiTModel": extract_minimax_h3_context,
|
|
# Future models:
|
|
# "CogVideoXTransformer3DModel": extract_cogvideox_context,
|
|
}
|
|
|
|
|
|
def register_extractor(transformer_cls_name: str, extractor_fn: Callable) -> None:
|
|
"""
|
|
Register a new extractor function for a model type.
|
|
|
|
This allows extending TeaCache support to new models without modifying
|
|
the core TeaCache code.
|
|
|
|
Args:
|
|
transformer_cls_name: Transformer model type identifier (class name or type string)
|
|
extractor_fn: Function with signature (module, *args, **kwargs) -> CacheContext
|
|
|
|
Example:
|
|
>>> def extract_flux_context(module, hidden_states, timestep, guidance=None, **kwargs):
|
|
... # Preprocessing
|
|
... temb = module.time_text_embed(timestep, guidance)
|
|
... # Extract modulated input
|
|
... modulated = module.transformer_blocks[0].norm1(hidden_states, emb=temb)
|
|
... # Define execution
|
|
... def run_blocks():
|
|
... h = hidden_states
|
|
... for block in module.transformer_blocks:
|
|
... h = block(h, temb=temb)
|
|
... return (h,)
|
|
... # Define postprocessing
|
|
... def postprocess(h):
|
|
... return module.proj_out(module.norm_out(h, temb))
|
|
... # Return context
|
|
... return CacheContext(modulated, hidden_states, None, temb, run_blocks, postprocess)
|
|
>>> register_extractor("FluxTransformer2DModel", extract_flux_context)
|
|
"""
|
|
EXTRACTOR_REGISTRY[transformer_cls_name] = extractor_fn
|
|
|
|
|
|
def get_extractor(transformer_cls_name: str) -> Callable:
|
|
"""
|
|
Get extractor function for given transformer class.
|
|
|
|
This function looks up the extractor based on the exact transformer_cls_name string,
|
|
which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__).
|
|
|
|
Args:
|
|
transformer_cls_name: Transformer class name (e.g., "QwenImageTransformer2DModel")
|
|
Must exactly match a key in EXTRACTOR_REGISTRY.
|
|
|
|
Returns:
|
|
Extractor function with signature (module, *args, **kwargs) -> CacheContext
|
|
|
|
Raises:
|
|
ValueError: If model type not found in registry
|
|
|
|
Example:
|
|
>>> # Get extractor for QwenImageTransformer2DModel
|
|
>>> extractor = get_extractor("QwenImageTransformer2DModel")
|
|
>>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...)
|
|
"""
|
|
# Direct lookup - no substring matching
|
|
if transformer_cls_name in EXTRACTOR_REGISTRY:
|
|
return EXTRACTOR_REGISTRY[transformer_cls_name]
|
|
|
|
# No match found
|
|
available_types = list(EXTRACTOR_REGISTRY.keys())
|
|
raise ValueError(
|
|
f"Unknown model type: '{transformer_cls_name}'. "
|
|
f"Available types: {available_types}\n"
|
|
f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY."
|
|
)
|