From edbd610060ee7a0dbf586da721bfaa16d375815d Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Tue, 23 Dec 2025 18:21:23 +0800 Subject: [PATCH] [RFC] Clean and refactor TeaCache registry (#416) Signed-off-by: yuanheng --- docs/user_guide/acceleration/teacache.md | 1 + tests/diffusion/cache/test_cache_backends.py | 3 + vllm_omni/diffusion/cache/teacache/backend.py | 16 ++--- vllm_omni/diffusion/cache/teacache/config.py | 28 +++++---- .../diffusion/cache/teacache/extractors.py | 48 +++++++-------- vllm_omni/diffusion/cache/teacache/hook.py | 16 +++-- .../qwen_image/pipeline_qwen_image_edit.py | 3 - .../pipeline_qwen_image_edit_plus.py | 1 - .../qwen_image/pipeline_qwen_image_layered.py | 60 ++++++++----------- 9 files changed, 89 insertions(+), 87 deletions(-) diff --git a/docs/user_guide/acceleration/teacache.md b/docs/user_guide/acceleration/teacache.md index e1a4db40c..8331c60f8 100644 --- a/docs/user_guide/acceleration/teacache.md +++ b/docs/user_guide/acceleration/teacache.md @@ -111,6 +111,7 @@ th { | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | +| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ### VideoGen diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 7bc49bdb7..ed9301410 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -138,6 +138,7 @@ class TestTeaCacheBackend: mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "QwenImagePipeline" mock_transformer = Mock() + mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel" mock_pipeline.transformer = mock_transformer config = DiffusionCacheConfig(rel_l1_thresh=0.3) @@ -154,6 +155,7 @@ class TestTeaCacheBackend: mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "QwenImagePipeline" mock_transformer = Mock() + mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel" mock_pipeline.transformer = mock_transformer config = DiffusionCacheConfig(rel_l1_thresh=0.3, coefficients=[1.0, 0.5, 0.2, 0.1, 0.05]) @@ -169,6 +171,7 @@ class TestTeaCacheBackend: mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "QwenImagePipeline" mock_transformer = Mock() + mock_transformer.__class__.__name__ = "QwenImageTransformer2DModel" mock_pipeline.transformer = mock_transformer # Mock hook registry diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 108dd5eb2..9db153cc4 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -47,20 +47,20 @@ class TeaCacheBackend(CacheBackend): and applies the TeaCache hook to the transformer. Args: - pipeline: Diffusion pipeline instance. Extracts transformer and model_type: + pipeline: Diffusion pipeline instance. Extracts transformer and transformer_type: - transformer: pipeline.transformer - - model_type: pipeline.__class__.__name__ + - transformer_type: pipeline.transformer.__class__.__name__ """ - # Extract transformer and model_type from pipeline + # Extract transformer and transformer_type from pipeline transformer = pipeline.transformer - model_type = pipeline.__class__.__name__ + transformer_type = transformer.__class__.__name__ - # Create TeaCacheConfig from DiffusionCacheConfig with model_type + # Create TeaCacheConfig from DiffusionCacheConfig with transformer_type # Access parameters via attribute access: config.rel_l1_thresh # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig try: teacache_config = TeaCacheConfig( - model_type=model_type, + transformer_type=transformer_type, rel_l1_thresh=self.config.rel_l1_thresh, coefficients=self.config.coefficients, ) @@ -69,7 +69,7 @@ class TeaCacheBackend(CacheBackend): raise ValueError( f"Invalid TeaCache configuration: {e}. " f"Expected keys: rel_l1_thresh, coefficients (optional). " - f"model_type is automatically extracted from pipeline.__class__.__name__." + f"transformer_type is automatically extracted from pipeline.transformer.__class__.__name__." ) # Apply hook to transformer @@ -80,7 +80,7 @@ class TeaCacheBackend(CacheBackend): logger.info( f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " - f"model_type={teacache_config.model_type}" + f"transformer_class={teacache_config.transformer_type}" ) def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 5a0cac626..916832bbf 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -7,17 +7,18 @@ from dataclasses import dataclass # These coefficients account for model-specific characteristics in how embeddings change # Source: TeaCache paper and ComfyUI-TeaCache empirical tuning _MODEL_COEFFICIENTS = { - # FLUX model coefficients from TeaCache paper - "FluxPipeline": [ + # FLUX transformer coefficients from TeaCache paper + "FluxTransformer2DModel": [ 4.98651651e02, -2.83781631e02, 5.58554382e01, -3.82021401e00, 2.64230861e-01, ], - # Qwen-Image model coefficients from ComfyUI-TeaCache + # Qwen-Image transformer coefficients from ComfyUI-TeaCache # Tuned specifically for Qwen's dual-stream transformer architecture - "QwenImagePipeline": [ + # Used for all Qwen-Image Family pipelines, in general + "QwenImageTransformer2DModel": [ -4.50000000e02, 2.80000000e02, -4.50000000e01, @@ -43,15 +44,15 @@ class TeaCacheConfig: - 0.4: ~1.8x speedup with slight quality loss - 0.6: ~2.0x speedup with noticeable quality loss coefficients: Polynomial coefficients for rescaling L1 distance. If None, uses - model-specific defaults based on model_type. - model_type: Pipeline class name (e.g., "QwenImagePipeline", "FluxPipeline"). - Must match OmniDiffusionConfig.model_class_name for proper extractor lookup. - Defaults to "QwenImagePipeline". + model-specific defaults based on transformer_type. + transformer_type: Transformer class name (e.g., "QwenImageTransformer2DModel"). + Auto-detected from pipeline.transformer.__class__.__name__ in backend. + Defaults to "QwenImageTransformer2DModel". """ rel_l1_thresh: float = 0.2 coefficients: list[float] | None = None - model_type: str = "QwenImagePipeline" + transformer_type: str = "QwenImageTransformer2DModel" def __post_init__(self) -> None: """Validate and set default coefficients.""" @@ -59,8 +60,13 @@ class TeaCacheConfig: raise ValueError(f"rel_l1_thresh must be positive, got {self.rel_l1_thresh}") if self.coefficients is None: - # Use model-specific coefficients, fallback to FluxPipeline if not found - self.coefficients = _MODEL_COEFFICIENTS.get(self.model_type, _MODEL_COEFFICIENTS["FluxPipeline"]) + # Use model-specific coefficients, explicitly check if the type exists or not + if self.transformer_type not in _MODEL_COEFFICIENTS: + raise KeyError( + f"Cannot find coefficients for {self.transformer_type}. " + f"Supported: {list(_MODEL_COEFFICIENTS.keys())}" + ) + self.coefficients = _MODEL_COEFFICIENTS[self.transformer_type] if len(self.coefficients) != 5: raise ValueError(f"coefficients must contain exactly 5 elements, got {len(self.coefficients)}") diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 8047e964e..56b4deefb 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -151,6 +151,7 @@ def extract_qwen_context( 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: @@ -170,6 +171,7 @@ def extract_qwen_context( 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 @@ -193,9 +195,9 @@ def extract_qwen_context( guidance = guidance.to(hidden_states.dtype) * 1000 temb = ( - module.time_text_embed(timestep, hidden_states) + module.time_text_embed(timestep, hidden_states, additional_t_cond) if guidance is None - else module.time_text_embed(timestep, guidance, hidden_states) + else module.time_text_embed(timestep, guidance, hidden_states, additional_t_cond) ) image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) @@ -254,22 +256,20 @@ def extract_qwen_context( # Registry for model-specific extractors -# Key: Pipeline class name (from OmniDiffusionConfig.model_class_name) +# Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext # -# Note: Use the exact pipeline class name as specified in OmniDiffusionConfig.model_class_name -# to ensure consistent lookup without substring matching. +# 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] = { - "QwenImagePipeline": extract_qwen_context, - "QwenImageEditPipeline": extract_qwen_context, - "QwenImageEditPlusPipeline": extract_qwen_context, + "QwenImageTransformer2DModel": extract_qwen_context, # Future models: - # "FluxPipeline": extract_flux_context, - # "CogVideoXPipeline": extract_cogvideox_context, + # "FluxTransformer2DModel": extract_flux_context, + # "CogVideoXTransformer3DModel": extract_cogvideox_context, } -def register_extractor(model_identifier: str, extractor_fn: Callable) -> None: +def register_extractor(transformer_cls_name: str, extractor_fn: Callable) -> None: """ Register a new extractor function for a model type. @@ -277,7 +277,7 @@ def register_extractor(model_identifier: str, extractor_fn: Callable) -> None: the core TeaCache code. Args: - model_identifier: Model type identifier (class name or type string) + transformer_cls_name: Transformer model type identifier (class name or type string) extractor_fn: Function with signature (module, *args, **kwargs) -> CacheContext Example: @@ -299,19 +299,19 @@ def register_extractor(model_identifier: str, extractor_fn: Callable) -> None: ... return CacheContext(modulated, hidden_states, None, temb, run_blocks, postprocess) >>> register_extractor("FluxTransformer2DModel", extract_flux_context) """ - EXTRACTOR_REGISTRY[model_identifier] = extractor_fn + EXTRACTOR_REGISTRY[transformer_cls_name] = extractor_fn -def get_extractor(model_type: str) -> Callable: +def get_extractor(transformer_cls_name: str) -> Callable: """ - Get extractor function for given model type. + Get extractor function for given transformer class. - This function looks up the extractor based on the exact model_type string, - which should match the pipeline class name from OmniDiffusionConfig.model_class_name. + 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: - model_type: Model type string (e.g., "QwenImagePipeline", "FluxPipeline") - Must exactly match a key in EXTRACTOR_REGISTRY. + 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 @@ -320,18 +320,18 @@ def get_extractor(model_type: str) -> Callable: ValueError: If model type not found in registry Example: - >>> # Get extractor for Qwen model - >>> extractor = get_extractor("QwenImagePipeline") + >>> # Get extractor for QwenImageTransformer2DModel + >>> extractor = get_extractor("QwenImageTransformer2DModel") >>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...) """ # Direct lookup - no substring matching - if model_type in EXTRACTOR_REGISTRY: - return EXTRACTOR_REGISTRY[model_type] + 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: '{model_type}'. " + 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." ) diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index b2f731665..1dfea9192 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/vllm_omni/diffusion/cache/teacache/hook.py @@ -62,7 +62,7 @@ class TeaCacheHook(ModelHook): def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: """ - Initialize hook with extractor from config model type. + Initialize hook with extractor from config transformer model type. Args: module: The module to initialize the hook for. @@ -70,9 +70,9 @@ class TeaCacheHook(ModelHook): Returns: The initialized module. """ - # Get extractor function based on model_type from config - # model_type should be the pipeline class name (e.g., "QwenImagePipeline") - self.extractor_fn = get_extractor(self.config.model_type) + # Get extractor function based on transformer_type from config + # transformer_type is the transformer class name (e.g., "QwenImageTransformer2DModel") + self.extractor_fn = get_extractor(self.config.transformer_type) # Set default context self.state_manager.set_context("teacache") @@ -244,9 +244,13 @@ def apply_teacache_hook(module: torch.nn.Module, config: TeaCacheConfig) -> None config: TeaCacheConfig specifying caching parameters Example: - >>> config = TeaCacheConfig(rel_l1_thresh=0.2, model_type="Qwen") + >>> config = TeaCacheConfig( + ... rel_l1_thresh=0.2, + ... transformer_type="QwenImageTransformer2DModel" + ... ) >>> apply_teacache_hook(transformer, config) - >>> # Model now uses TeaCache automatically, no code changes needed! + >>> # Transformer bound to the pipeline now uses TeaCache automatically, + ... # no code changes needed! """ registry = HookRegistry.get_or_create(module) hook = TeaCacheHook(config) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 8d180e625..e6374763c 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -230,9 +230,6 @@ class QwenImageEditPipeline( model, subfolder="processor", local_files_only=local_files_only ) - # Initialize cache backend to None (will be set by worker if needed) - self._cache_backend = None - self.stage = None self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 6068cf0f1..2970987f1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -194,7 +194,6 @@ class QwenImageEditPlusPipeline( ) self.stage = None - self._cache_backend = None self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index 45a165b7a..33ceeddbd 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -264,8 +264,6 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as ) self.default_sample_size = 128 - # VLLM-omni related, Initialize cache backend to None (will be set by worker if needed) - self._cache_backend = None self.stage = None def check_inputs( @@ -570,42 +568,36 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - # Forward pass for positive prompt (or unconditional if no CFG) - # cache_branch is passed to hook for CFG-aware state management - transformer_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": prompt_embeds_mask, - "encoder_hidden_states": prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": txt_seq_lens, - "attention_kwargs": self.attention_kwargs, - "additional_t_cond": is_rgb, - "return_dict": False, - } - if self._cache_backend is not None: - transformer_kwargs["cache_branch"] = "positive" - noise_pred = self.transformer(**transformer_kwargs)[0] + self.transformer.do_true_cfg = do_true_cfg # used in teacache hook + # Forward pass for positive prompt (or unconditional if no CFG) + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=prompt_embeds_mask, + encoder_hidden_states=prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + attention_kwargs=self.attention_kwargs, + additional_t_cond=is_rgb, + return_dict=False, + )[0] noise_pred = noise_pred[:, : latents.size(1)] if do_true_cfg: - neg_transformer_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": negative_prompt_embeds_mask, - "encoder_hidden_states": negative_prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": negative_txt_seq_lens, - "attention_kwargs": self.attention_kwargs, - "additional_t_cond": is_rgb, - "return_dict": False, - } - if self._cache_backend is not None: - neg_transformer_kwargs["cache_branch"] = "negative" - neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0] + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=negative_prompt_embeds_mask, + encoder_hidden_states=negative_prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=negative_txt_seq_lens, + attention_kwargs=self.attention_kwargs, + additional_t_cond=is_rgb, + return_dict=False, + )[0] neg_noise_pred = neg_noise_pred[:, : latents.size(1)] comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)