[feature] stable_audio_open_1 teacache support (#1314)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

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:

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.

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,
+    ],
 }
 
 

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,