smajji/DhVaani-0.5-TensorRT-aarch64

DhVaani-0.5 — TensorRT/ONNX artefacts (aarch64 / Jetson)

ONNX exports and TensorRT 10.13 engines built on a Jetson AGX Orin (aarch64, CUDA 12.1) for DhVaani, the 27-language zero-shot Indic TTS fine-tuned from ZipVoice, plus the full export→build→verify pipeline code.

Contents

path	what
trt_engines/fm_decoder_trt_t{T}.plan	fp16 TensorRT engines — fixed sequence-length ladder T ∈ {256, 512, 768, 1024, 1536, 2048, 2560} (each ≈ 256 MB)
onnx_models/fm_decoder[_trt].onnx	flow-matching decoder ONNX (≈ 456 MB; _trt = rank-1-input variant used to build engines)
onnx_models/text_encoder[_trt].onnx	text encoder ONNX (≈ 18 MB)
scripts/	full pipeline: export_onnx.py, build_trt.sh, check_onnx.py, infer_dwani.py (pytorch/onnx/trt backends), compare_audio.py, analyze_case.py, scaling_converter.py
run_dwani_trt.sh	one-command pipeline (venv → export → check → build → verify)
requirements.txt	python deps

Model I/O (fm_decoder engine)

One engine call = one Euler step for a chunk, running both classifier-free-guidance branches (batch=2 inside the graph) and returning the guidance-combined velocity:

• t(1), x(1,T,100), text_condition(1,T,100), speech_condition(1,T,100), padding_mask(1,T, bool, True=padded), guidance_scale(1) → v(1,T,100)
• The padding_mask input is what makes fixed-shape engines exact: padded frames are excluded from attention, so the real content is numerically identical to an unpadded run.

Usage (C++ / Jetson)

// pick smallest engine with T >= your chunk length (prompt+text frames)
//   T frames @ 24 kHz, hop 256 = 93.75 frames per audio second
// pad x/text/speech to T (zeros), set padding_mask[.., realT:] = true
auto engine = loadEngine("fm_decoder_trt_t768.plan");   // TRT C++ API
// per step: compute v, x += v * dt  (16 steps default)

The .plan files load directly with the TensorRT C++ API — no Python needed. Engines are architecture-locked: rebuild with scripts/build_trt.sh on the target machine from the ONNX files (works on x86 and aarch64, TRT 10.x).

Benchmarks (measured on this Jetson, AGX Orin, fp16)

engine T	~audio	GPU compute / call (one step)
256	2.7 s	11.7 ms
512	5.5 s	19.4 ms
768	8.2 s	29.2 ms
1024	10.9 s	40.0 ms
1536	16.4 s	68.9 ms
2048	21.8 s	103.9 ms
2560	27.3 s	146.9 ms

→ 16 steps, 2.9 s chunk @ T768: ≈ 0.47 s sampling (acoustic RTF ≈ 0.16, 6× real-time); 8 steps → RTF ≈ 0.08. (Latency is per call incl. both CFG branches.)

Audio quality: TRT fp16 @ 16 steps is audibly identical to PyTorch fp32 (sample comparisons in the pipeline: mel err ≈ 0.01–0.02, waveform corr ≈ 0.97–0.99 across ta/hi/kn). see scripts/compare_audio.py.

Notes on dynamic shapes

Dynamic-shape builds crash TRT 10.13's "compiler backend" for this graph on both x86 (CUDA-13 driver) and this aarch64 Jetson — hence the fixed-T ladder. The --nvinfer build recipe is in scripts/build_trt.sh; TRT_SIZES=dynamic is provided to retry a single dynamic engine on other TRT versions.

Licence & credit

Apache-2.0 (see LICENSE). Model: ARTPARK-IISc/DhVaani-0.5, based on k2-fsa/ZipVoice.