[Bounty $1500] AudioX bring up using TTNN APIs

Bounties
1

Tenstorrent Bounties

Issue ID: I_kwDOI9Wqc87XHo88

:memo: Background

This bounty is for bringing up the AudioX model using TTNN APIs on Tenstorrent hardware (Wormhole or Blackhole).

AudioX is a unified framework for anything-to-audio generation from HKUST, supporting multimodal control signals. Released in 2025, key features include:

  • Anything-to-audio: Supports text, video, image, and audio inputs for audio generation
  • Multiple tasks: Text-to-audio, text-to-music, video-to-audio, video-to-music, audio inpainting, music completion
  • Multimodal Adaptive Fusion: Novel module for effective fusion of diverse multimodal inputs
  • Large-scale training: IF-caps dataset with 7 million high-quality samples
  • State-of-the-art performance: Superior results on AudioCaps, MusicCaps, T2A-bench, and AudioTime benchmarks
  • Diffusion Transformer: Uses diffusion-based generation with transformer architecture
  • Instruction following: Handles complex temporal instructions (timestamps, ordering, counting)
  • High quality: 16 kHz audio output with natural sound quality

AudioX achieves superior performance across benchmarks, especially in text-to-audio and text-to-music generation, demonstrating powerful multimodal control and instruction-following capabilities.

The goal is to enable this model to run on Tenstorrent hardware for high-throughput, low-latency multimodal audio generation .

:bullseye: What Success Looks Like

A successful submission will fulfill all requirements in the following stages. Payout is made after all three stages are completed.

Stage 1 — Bring-Up

  • Implement AudioX using TTNN APIs (Python)
  • Implements the full generation pipeline:
    • Multimodal encoders (text, video, image, audio)
    • Multimodal Adaptive Fusion module
    • Diffusion Transformer decoder
    • Vocoder for waveform generation
  • Model runs on either N150 or N300 Tenstorrent hardware with no errors
  • Supports multiple generation modes:
    • Text-to-audio: Generate sound effects from text descriptions
    • Text-to-music: Generate music from text prompts
    • Video-to-audio: Generate audio synchronized with video
    • Video-to-music: Generate background music for video
  • Produces valid 16 kHz audio output on sample inputs
  • Output is verifiable (audio quality assessment, compare with PyTorch reference)
  • Achieves baseline throughput target:
    • At least 20 tokens/second for diffusion sampling
    • Generation time < 30 seconds for 10-second audio
  • Accuracy evaluation:
    • Inception Score comparable to paper results
    • Text-audio alignment score > 0.7
    • Token-level accuracy > 95% against PyTorch reference
  • Audio quality: Natural sound quality and accurate multimodal alignment
  • Clear instructions for setup and running the model

Stage 2 — Basic Optimizations

  • Use optimal sharded/interleaved memory configs for transformer layers
  • Implement efficient sharding strategy for:
    • Text encoder (CLAP or similar)
    • Video encoder (if applicable)
    • Multimodal Adaptive Fusion module
    • Diffusion Transformer layers
    • Attention mechanisms (self and cross-attention)
  • Fuse simple ops where possible (e.g., layer normalization, attention patterns, activation functions)
  • Store intermediate activations in L1 where beneficial
  • Use recommended TTNN/tt-metal diffusion model flows
  • Leverage TT library of fused ops for attention and MLP blocks
  • Optimize multimodal fusion operations
  • Efficient diffusion sampling (DDPM/DDIM)
  • Optimize vocoder integration

Stage 3 — Deeper Optimization

  • Maximize core counts used per inference
  • Implement deeper TT-specific optimizations:
    • Efficient diffusion sampling with fewer steps (DDIM, DPM-Solver)
    • Flash Attention or equivalent for attention layers
    • Minimize generation latency
    • Batch processing for multiple audio generations
    • Pipeline multimodal encoding with diffusion sampling
    • Optimize cross-modal attention
    • Efficient temporal instruction handling
  • Minimize memory and TM (tensor manipulation) overheads
  • Explore classifier-free guidance optimization
  • Document any advanced tuning, known limitations, or trade-offs
  • Target stretched goals:
    • 50+ tokens/second diffusion sampling
    • Generation time < 10 seconds for 10-second audio
    • Support for longer audio (30+ seconds)
    • Efficient multi-task switching

:compass: Guidance & Starting Points

  • Use the TTNN model bring-up tech report as your primary reference
  • Reference diffusion model implementations in tt-metal for diffusion patterns
  • Reference transformer implementations in tt-metal for transformer patterns
  • Use the AudioX project page for demos and examples
  • Refer to the AudioX paper (arXiv:2503.10522) for technical details
  • Check AudioX HuggingFace for model checkpoints
  • Check AudioX GitHub for official code repository
  • Refer to TT Fused ops PR #29236 for optimization opportunities
  • The model architecture consists of:
    • Multimodal encoders: Separate encoders for text, video, image, audio inputs
    • Multimodal Adaptive Fusion: Fuses diverse modalities with cross-modal alignment
    • Diffusion Transformer: Denoising network for audio generation
    • Vocoder: Converts latent representations to waveforms
    • Classifier-free guidance: For controllable generation
  • Key challenges:
    • Multimodal fusion optimization
    • Diffusion sampling efficiency (multiple denoising steps)
    • Cross-modal attention across different input types
    • Temporal instruction handling (timestamps, ordering, counting)
    • Vocoder integration
    • Long audio generation
  • Ask for help or file issues if ops are missing in TTNN

:magnifying_glass_tilted_right: Possible Approaches

  • Start from the official AudioX implementation and port components sequentially:
    1. Text encoder (CLAP or similar)
    2. Video/image encoders (if using video-to-audio)
    3. Multimodal Adaptive Fusion module
    4. Diffusion Transformer (U-Net or DiT architecture)
    5. Vocoder integration
    6. Diffusion sampling loop (DDPM/DDIM)
    7. Classifier-free guidance
  • Validate each submodule’s output against PyTorch reference before integration
  • Experiment with different sharding strategies:
    • Attention head sharding
    • Sequence length sharding
    • Batch dimension sharding for diffusion steps
  • Use TTNN profiling tools to identify bottlenecks in:
    • Multimodal encoding time
    • Fusion module overhead
    • Diffusion sampling iterations
    • Cross-modal attention
    • Vocoder latency
  • Test diverse use cases:
    • Text-to-audio with complex descriptions
    • Text-to-music with style instructions
    • Video-to-audio synchronization
    • Temporal instructions (timestamps, ordering)
    • Audio inpainting and music completion
  • Open a draft PR early to get feedback on your approach

:bar_chart: Result Submission Guidelines

Beyond the model implementation itself, contributors must submit the following material as proof of work.

However, feel free to open a PR at any time if you want us checking that you are on the right track.

Just understand that payout is only made after all 3 stages are completed.

Deliverables:

  • Functional model implementation in TTNN
  • Validation logs showing output correctness (latent accuracy vs PyTorch)
  • Performance report with generation time and quality metrics
  • Sample audio outputs comparing TTNN vs PyTorch:
    • Text-to-audio examples (sound effects)
    • Text-to-music examples (various genres)
    • Video-to-audio examples (if implemented)
    • Temporal instruction examples (timestamps, ordering)
  • Inception Score and other quality metrics
  • Performance header for final review

Links:

:books: Resources

Model Resources

Benchmark Resources

TT-Metal Resources