Add Apple Silicon (MPS) and CPU support

[Talpik]

Runs FireRedTTS3 on Apple Silicon (MPS) and on CPU, without changing anything on CUDA.

Includes the one-line fix from #5 (first commit), since the Instruct cloning path is one of the things this PR tests. If #5 lands first, this rebases cleanly.

What was CUDA-only

Device, compute dtype, attention implementation and FFT capability are now resolved once in a new fireredtts3/utils/device.py, so the model code stays free of per-backend branches:

Replaced	With
torch.device('cuda') (base + instruct)	get_device() — CUDA → MPS → CPU
@torch.autocast(device_type='cuda', dtype=bfloat16) (×3)	@autocast
@autocast('cuda', enabled=False) (×2, rotary)	@disable_autocast
attn_implementation='flash_attention_2' (×3)	get_attn_implementation()
unguarded torch.cuda.manual_seed*	guarded, plus torch.mps.manual_seed

flash_attn has no macOS wheel at all, so off CUDA this falls back to PyTorch SDPA (both core models already declare _supports_sdpa = True).

Backends without FFT / complex kernels — MPS on older torch — additionally need the RedAE ISTFT and the CAM++ kaldi fbank on CPU. That is detected by a runtime probe rather than a version check, so it stays dormant where MPS does support FFT (torch 2.8 on macOS 26 does) and switches on automatically where it does not.

Memory: FIRERED_WEIGHT_DTYPE

Autocast only casts activations, so it cannot shrink resident memory — relevant when 3.06B fp32 params do not fit. This adds a knob to load the checkpoints in half precision. Making it work required the dtype to follow the weights rather than being pinned to fp32 at the pipeline entry points (prompt latents, speaker embedding, flow-matching noise x0, t_span, empty latents_out); without that, MPS aborts with

failed assertion `Destination NDArray and Accumulator NDArray cannot have different datatype in MPSNDArrayMatrixMultiplication'

The ISTFT deliberately stays in fp32, since complex half is not a usable dtype.

Measured on an M4 Max (macOS 26, torch 2.8), Base zero-shot cloning of a 3-sentence paragraph over two reference voices, RTF as the range across 3 interleaved runs, speaker similarity via the bundled CAM++:

Weights	Resident weights	RTF	Speaker sim
float32 (default)	11.42 GiB	0.87–0.96	0.947 / 0.935
bfloat16	5.71 GiB	0.73–0.76	0.941 / 0.933

On MPS bfloat16 weights are the setting to want: half the memory and ~20% faster, at effectively unchanged similarity. The default stays float32 because it matches the checkpoints and leaves CUDA untouched.

Correction: an earlier revision of this description reported bf16 as slower (RTF 1.47 vs 1.15) from a table timing a single short sentence, where fixed per-call overhead dominates. Re-measured on a realistic workload, 3 interleaved runs each, that conclusion was wrong; the README commit in this PR carries the corrected numbers.

Autocast, by contrast, buys nothing on MPS (RTF 0.88–0.93 vs 0.85–0.91 over fp32 weights) — the DiT flow head, the real bottleneck, sits outside the autocast region. So autocast off is the default off CUDA, and CUDA keeps its bf16 autocast defaults unchanged.

Dependencies

requirements-mps.txt is the macOS set: no flash_attn, plus soundfile. torchcodec==0.7.0 refuses to load against FFmpeg 8 (it supports 4–7), and without it torchaudio has no I/O backend at all on macOS — even torchaudio.save fails.

Testing

All five inference paths, 24 kHz output, Russian reference voices (12.6–17.9 s prompts, ~10 s targets):

Path	Result
Base zero-shot cloning, 4 voices	RTF 0.81–0.90
Instruct ICL cloning	RTF 0.86
Voice design	RTF 0.84, coherent attribute plan
Acoustic edit, speed 0.8x	12.6 s → 15.68 s, i.e. exactly 1/0.8
Semantic edit	RTF 0.93; accurate on English, and the <|edit|> mask lands on the intended word

Speaker similarity (bundled CAM++, matched pairs on the diagonal) — the fp32 numbers are unchanged before and after the dtype refactor:

generated            lu-lo   lu-vei-in     lun-han  sun-da-dao   <- reference
lu-lo                0.947       0.842       0.810       0.836
lu-vei-in            0.841       0.935       0.773       0.840
lun-han              0.828       0.785       0.948       0.850
sun-da-dao           0.807       0.797       0.823       0.929

Note that MPS is not bit-reproducible run to run: two runs of identical code produce different bytes, so equivalence is established via durations and similarity scores rather than checksums.

Not tested on CUDA — I have no CUDA machine. The CUDA path is intended to be byte-identical to before: get_device() returns cuda first, get_attn_implementation() returns flash_attention_2, autocast defaults to bf16, FIRERED_WEIGHT_DTYPE defaults to the checkpoint's fp32, and the FFT probe passes so no CPU fallback engages. A confirmation run on CUDA would be welcome.

scripts/convert_to_bf16.py is a user-facing tool the README points at, not test code: it halves the checkpoints on disk (19 GB → 9.7 GB), verifying every tensor against the fp32 original for exact bfloat16 rounding before replacing it, so the two routes to half precision — cast on load via FIRERED_WEIGHT_DTYPE, or convert once and load with no variable set — are both documented and supported.

scripts/mps_smoke_test.py and scripts/spk_similarity.py, by contrast, are just the harness used above; happy to drop those two if you would rather keep the repo lean.

🤖 Generated with Claude Code