Irodori-TTS-burn

A full-scratch Rust reimplementation of Irodori-TTS using the Burn machine learning framework.

Irodori-TTS is a Japanese text-to-speech model built around a Diffusion Transformer (DiT) with RoPE positional encoding, Joint attention with KV caching, and a DACVAE neural codec. This project ports the inference pipeline and LoRA fine-tuning from PyTorch to Burn, preserving numerical parity with the Python reference.

Features

• Full-model port — DiT blocks, RoPE, Joint attention, DACVAE codec (15.3K LOC, 46 source files)
• Multi-backend — LibTorch f32/bf16/MPS, CUDA f32 (CubeCL), WGPU f32/f16, WgpuRaw f16; enum-based runtime dispatch
• KV cache — cached forward pass for the RF sampler loop
• LoRA — adapter merging for inference + LoRA fine-tuning with flow-matching loss
• Training — dataset loading, LR schedules, gradient clipping, atomic checkpointing, resume
• Numerical parity — layer-by-layer validation against Python fixtures (all PASS); E2E 4-step CFG verified on NdArray (exact), LibTorch MPS (exact), WgpuRaw f16 (5.29e-4, f16 rounding)
• E2E pipeline — pipeline binary: text → WAV
• Cargo feature flags — inference, codec, text-normalization, lora, train, cli
• 306 unit tests, clippy clean

Performance (RTX A6000)

Rust LibTorch bf16 is 63% faster than Python — and Python can't even run bf16. Rust LibTorch f32 is within 2.4% of Python — near-parity (FFI abstraction overhead).

Performance (M4 Pro Mac Mini)

RTF < 1 = faster than real-time. Joint CFG requires equal guidance scales (quality tradeoff). Alternating CFG alternates which signal is unconditioned each step — same speed as Joint (both batch=2). MPS benchmarks use LibTorch from Aratako/Irodori-TTS PyTorch venv (torch 2.10.0). WgpuRaw f16 is the recommended no-dependency fallback (avoids burn-fusion DACVAE crash). LibTorch MPS f16 is the recommended backend on Apple Silicon (1.66× faster than WgpuRaw f16).

Training throughput (LoRA fine-tuning, RTX A6000):

Cargo features

Backend selection is handled at runtime via --backend CLI flag. See InferenceBackendKind and TrainingBackendKind enums in src/backend_config.rs.

Quick start

See docs/setup.md for the full setup guide including:

• Setting up the Python reference environment
• Downloading and converting model weights
• Configuring the LibTorch backend
• Running validation and benchmarks

# Inference — WgpuRaw f16 (recommended no-dependency backend on Mac/Linux/Windows)
cargo run --release --features cli --bin pipeline -- \
    --backend wgpu-raw-f16 \
    --checkpoint target/model_converted.safetensors \
    --codec-weights target/dacvae_weights.safetensors \
    --text "こんにちは、世界！" \
    --output /tmp/out.wav

# Inference — LibTorch MPS (fastest on Apple Silicon; requires PyTorch venv)
cargo run --release --features cli --bin pipeline -- \
    --backend libtorch-mps \
    --checkpoint target/model_converted.safetensors \
    --codec-weights target/dacvae_weights.safetensors \
    --text "こんにちは、世界！" \
    --output /tmp/out.wav

# With LoRA adapter
cargo run --release --features "lora,cli" --bin pipeline -- \
    --checkpoint target/model_converted.safetensors \
    --codec-weights target/dacvae_weights.safetensors \
    --adapter path/to/adapter_dir \
    --text "こんにちは、世界！" \
    --output /tmp/out.wav

# LoRA training
cargo run --release --features "train,cli" --bin train_lora -- \
    --backend libtorch \
    --checkpoint target/model_converted.safetensors \
    --manifest path/to/manifest.jsonl

Repository structure

src/
├── lib.rs                  — public API (feature-gated modules)
├── model/                  — DiT, attention, RoPE, positional encoding
├── codec/                  — DACVAE encoder/decoder
├── rf.rs                   — RF (rectified-flow) sampler + KV cache
├── inference.rs            — InferenceBuilder type-state API
├── weights.rs              — weight loading (safetensors)
├── lora.rs                 — LoRA adapter merging
├── train/                  — LoRA training (dataset, trainer, loss, LR schedule)
├── text_normalization.rs   — Japanese text normaliser
├── config.rs               — model/training/sampling configuration
├── error.rs                — thiserror-based error types
└── bin/
    ├── pipeline.rs         — E2E inference binary
    ├── train_lora.rs       — LoRA fine-tuning binary
    ├── infer.rs            — inference-only binary
    ├── bench_realmodel.rs  — multi-backend latency benchmark
    └── ...
docs/
├── setup.md                — installation guide (start here)
├── feature-coverage.md     — implementation status + test tracking
├── benchmarks.md           — benchmark results
└── planning/               — architecture notes

Development

just ci           # fmt-check + clippy -D warnings + test
just validate     # numerical parity (requires ../Irodori-TTS venv)
just bench-tch    # LibTorch f32 benchmark
just quality-compare   # multi-backend WAV comparison

License

MIT — see LICENSE.