# TT-DiT: Tenstorrent Diffusion Transformers

This directory contains the implementation of the Tenstorrent Diffusion Transformer (TT-DiT) architecture, designed for optimized parallelism of DiTs for image and video generation.

## Supported Models

For detailed information about each model including performance metrics, usage instructions, and specific requirements, see:

- **[Stable Diffusion 3.5 Large](models/StableDiffusion35.md)** - Text-to-image generation
- **[Flux 1](models/Flux1.md)** - Text-to-image generation (schnell & dev variants)
- **[Motif](models/Motif.md)** - Text-to-image generation model
- **[Qwen-Image](models/QwenImage.md)** - Text-to-image generation model
- **[Ideogram 4.0](models/Ideogram4.md)** - Text-to-image generation model (best-in-class typography; Blackhole)
- **[Mochi-1](models/Mochi_1.md)** - Video generation model
- **[Wan2.2-T2V-A14B](models/Wan2_2.md)** - Text-to-video generation model
- See **[experimental/](experimental/)** for in-progress model variants (AniSora, Wan2.2-Distill, LoRA adapters)
- **[LTX-2.3](models/LTX2.md)** - Joint audio-video generation model (Pro and Fast variants)

## Directory Structure

```
tt_dit/
├── layers/              # Core neural network layers
├── models/              # Model architectures and documentation
│   ├── transformers/    # Transformer implementations (SD35, Mochi, Wan, LTX, Flux1, Motif, QwenImage)
│   ├── vae/            # VAE/Autoencoder implementations
│   ├── StableDiffusion35.md  # SD3.5 model documentation
│   ├── Flux1.md         # Flux 1 model documentation
│   ├── Motif.md         # Motif model documentation
│   ├── QwenImage.md     # Qwen-Image model documentation
│   ├── Mochi_1.md       # Mochi-1 model documentation
│   ├── Wan2_2.md        # Wan2.2 model documentation
│   └── LTX2.md          # LTX-2.3 model documentation
├── encoders/            # Text encoder implementations
│   ├── clip/           # CLIP encoder
│   ├── gemma/          # Gemma encoder
│   └── t5/             # T5 encoder
├── parallel/            # Parallelization utilities
│   ├── config.py        # Parallel configuration
│   └── manager.py       # Parallel execution management
├── pipelines/           # End-to-end model pipelines
│   ├── stable_diffusion_35_large/
│   ├── mochi/
│   ├── wan/
│   ├── ltx/
│   ├── flux1/
│   ├── motif/
│   └── qwenimage/
├── tests/              # Test suite
│   ├── models/         # Model-level tests (sd35, mochi, wan2_2, ltx, flux1, motif, qwenimage)
│   ├── encoders/       # Encoder tests
│   ├── blocks/         # Block-level tests
│   └── unit/          # Unit tests for layers
└── utils/             # Utility functions
    ├── check.py       # Validation utilities
    ├── padding.py     # Padding operations
    ├── substate.py    # State management
    └── tensor.py      # Tensor operations
```

## Core Components

### Layers
- **Embeddings**: Position and token embedding implementations
- **Feedforward**: MLP and feed-forward network components
- **Linear**: Optimized linear transformation layers
- **Normalization**: Layer normalization implementations

### Models
- **Transformers**: DiT transformer implementations for various generative models
  - Support for multiple architectures including SD3.5, Mochi, Wan2.2, LTX-2, Flux1, Motif, and Qwen-Image
  - Model-specific attention mechanisms and transformer architectures
- **Autoencoders**: VAE implementations for different models

### Parallel Processing
- **Config**: Parallel configuration management with `DiTParallelConfig`
- **Manager**: Parallel execution and device management

### Pipelines
End-to-end pipeline implementations for multiple generative models:
- **Stable Diffusion 3.5 Large**: Text-to-image generation (1024x1024px)
- **Flux 1**: Text-to-image generation (schnell & dev variants, 1024x1024px)
- **Motif**: Text-to-image generation (6B model, 1024x1024px)
- **Qwen-Image**: Text-to-image generation (1024x1024px)
- **Mochi-1**: Video generation model (824x480px, 168 frames)
- **Wan2.2-T2V-A14B**: Text-to-video generation
- **LTX-2.3**: Joint audio-video generation (Pro one-stage and Fast distilled two-stage)

Each pipeline includes:
- Automatic parallel configuration for different device meshes
- Optimized execution on Wormhole systems
- Comprehensive timing collection and profiling

## Testing

The test suite is organized into three main categories:

1. **Unit Tests** (`tests/unit/`):
   - Layer-level testing
   - Individual component validation
   - Coverage for embeddings, feedforward, linear, and normalization layers

2. **Model Tests** (`tests/models/`):
   - Smoke tests for model components
   - Performance testing
   - Pipeline validation

Running tests:
```bash
# Run tests from tt-metal/

# Run unit tests
python -m pytest  models/tt_dit/tests/unit/

# Run all model tests
python -m pytest  models/tt_dit/tests/models/

# Run specific model pipeline tests
python -m pytest models/tt_dit/tests/models/sd35/test_pipeline_sd35.py -v
python -m pytest models/tt_dit/tests/models/flux1/test_pipeline_flux1.py -v
python -m pytest models/tt_dit/tests/models/motif/test_pipeline_motif.py -v
python -m pytest models/tt_dit/tests/models/qwenimage/test_pipeline_qwenimage.py -v
python -m pytest models/tt_dit/tests/models/mochi/test_pipeline_mochi.py -v
python -m pytest models/tt_dit/tests/models/wan2_2/test_pipeline_wan.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py -v
```

## Key Features

- **Modular Architecture**: Clean separation between layers, models, and parallel processing
- **Optimized Parallelism**:
  - Automatic configuration of tensor and sequence parallelism
  - Efficient device mesh management
  - Dedicated parallel managers for different model components
- **Performance Monitoring**:
  - Built-in timing collection
  - Performance profiling support
  - Tracing capabilities for optimization

## Contributing

When adding new components:
1. Place implementation in appropriate subdirectory
2. Add corresponding unit tests
3. Update model tests if necessary
4. Ensure parallel processing compatibility