diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 11836bbd7..5edfa2237 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -93,6 +93,12 @@ The following table shows which models are currently supported by each accelerat | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | +### AudioGen + +| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | +|------------------------|------------------------------------------|:--------:|:---------:|:----------:|:--------------:|:------------:| +| **Stable-Audio-Open** | `stabilityai/stable-audio-open-1.0` | ✅ | | | | | + ## Performance Benchmarks The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps: diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index cdacf49dc..f3a278b21 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py @@ -7,12 +7,16 @@ from typing import Any import numpy as np import torch from vllm.config import LoadConfig +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.cache.teacache.extractors import get_extractor from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline +from vllm_omni.diffusion.models.stable_audio.pipeline_stable_audio import StableAudioPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams class DataCollectionHook(ModelHook): @@ -32,16 +36,15 @@ class DataCollectionHook(ModelHook): def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any: ctx = self.extractor_fn(module, *args, **kwargs) - modulated_input_cpu = ctx.modulated_input.detach().cpu().numpy() + modulated_input_cpu = ctx.modulated_input.detach().float().cpu().numpy() outputs = ctx.run_transformer_blocks() ctx.hidden_states = outputs[0] if len(outputs) > 1 and ctx.encoder_hidden_states is not None: ctx.encoder_hidden_states = outputs[1] - model_output_cpu = ctx.hidden_states.detach().cpu().numpy() + model_output_cpu = ctx.hidden_states.detach().float().cpu().numpy() self.current_trajectory.append((modulated_input_cpu, model_output_cpu)) - return ctx.postprocess(ctx.hidden_states) def start_collection(self): @@ -82,6 +85,32 @@ class BagelAdapter: transformer._forward_flow = transformer.forward +class StableAudioAdapter: + """Adapter for Stable Audio Open 1.0 coefficient estimation.""" + + @staticmethod + def load_pipeline(model_path: str, device: str = "cuda", dtype: torch.dtype = torch.float16) -> Any: + od_config = OmniDiffusionConfig.from_kwargs(model=model_path, dtype=dtype) + + # Strictly necessary because we bypass loader.load_model() + with set_default_torch_dtype(dtype): + pipeline = StableAudioPipeline(od_config=od_config) + + loader = DiffusersPipelineLoader(LoadConfig()) + loader.load_weights(pipeline) + pipeline.to(device) + return pipeline + + @staticmethod + def get_transformer(pipeline: Any) -> tuple[Any, str]: + return pipeline.transformer, "StableAudioDiTModel" + + @staticmethod + def install_hook(transformer: Any, hook: DataCollectionHook) -> None: + registry = HookRegistry.get_or_create(transformer) + registry.register_hook(hook._HOOK_NAME, hook) + + class DefaultAdapter: """Default adapter for standard diffusers pipelines.""" @@ -101,6 +130,7 @@ class DefaultAdapter: _MODEL_ADAPTERS: dict[str, type] = { "Bagel": BagelAdapter, + "StableAudio": StableAudioAdapter, } _EPSILON = 1e-6 @@ -165,17 +195,19 @@ class TeaCacheCoefficientEstimator: def collect_from_prompt(self, prompt: str, **generate_kwargs): self.hook.start_collection() - from vllm_omni.diffusion.request import OmniDiffusionRequest - req = OmniDiffusionRequest( - prompt=prompt, - num_inference_steps=generate_kwargs.get("num_inference_steps", 20), - seed=generate_kwargs.get("seed", 42), + prompts=[prompt], + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=generate_kwargs.get("num_inference_steps", 20), + seed=generate_kwargs.get("seed", 42), + ), ) - self.pipeline.forward(req) + with torch.no_grad(): + self.pipeline.forward(req) trajectory = self.hook.stop_collection() if trajectory: self.collected_data.append(trajectory) + torch.cuda.empty_cache() def estimate(self, poly_order: int = 4) -> list[float]: """Estimate polynomial coefficients from collected data. diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 9fc7dc5b1..8568ad660 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -46,6 +46,14 @@ _MODEL_COEFFICIENTS = { 3.20000000e00, -2.00000000e-02, ], + # Estimated TeaCache polynomial coefficients for StableAudioDiTModel. + "StableAudioDiTModel": [ + 121.77490545701518, + -153.7449426160371, + 68.05368574596551, + -12.281286412689623, + 1.0733905006198015, + ], } diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 56ff65b7c..bdb3f6a78 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -721,6 +721,112 @@ def extract_flux2_klein_context( ) +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, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -732,6 +838,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = { "Bagel": extract_bagel_context, "ZImageTransformer2DModel": extract_zimage_context, "Flux2Klein": extract_flux2_klein_context, + "StableAudioDiTModel": extract_stable_audio_context, # Future models: # "FluxTransformer2DModel": extract_flux_context, # "CogVideoXTransformer3DModel": extract_cogvideox_context,