Add Apple Silicon (MPS) and CPU support

.gitignore

@@ -14,4 +14,7 @@ pretrained_models/
 
 # Eval
 data
-data/
+data/
+# Local venv / test output
+.venv/
+outputs/

README.md

@@ -79,6 +79,74 @@
 pip install -r requirements.txt
 ```
 
+### Installation on Apple Silicon (MPS)
+
+`flash_attn` has no macOS wheel, and `torchcodec` needs FFmpeg <= 7, so use the
+macOS requirement set instead (attention falls back to PyTorch SDPA, audio I/O
+to `soundfile`):
+
+```sh
+uv venv --python 3.12 .venv
+uv pip install --python .venv/bin/python -r requirements-mps.txt
+```
+
+The device and compute dtype are resolved once in `fireredtts3/utils/device.py`
+(CUDA -> MPS -> CPU) and can be overridden:
+
+| Variable | Values | Default |
+| --- | --- | --- |
+| `FIRERED_DEVICE` | `cuda` / `mps` / `cpu` | first available |
+| `FIRERED_DTYPE` | autocast dtype (activations) | `bfloat16` on CUDA, `float32` elsewhere |
+| `FIRERED_WEIGHT_DTYPE` | dtype the checkpoints are loaded in | as stored (fp32 officially) |
+
+`FIRERED_DTYPE` only casts activations, so it does **not** shrink resident
+memory; `FIRERED_WEIGHT_DTYPE` is the knob that does. Measured on an M4 Max
+(torch 2.8, macOS 26), Base zero-shot cloning of a 3-sentence paragraph over
+two reference voices, RTF as the range across 3 interleaved runs:
+
+| Weights | Resident weights | RTF | Speaker sim |
+| --- | --- | --- | --- |
+| `float32` (official checkpoints) | 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 what you want:** half the memory *and* ~20%
+faster, at effectively unchanged speaker similarity. There are two ways to get
+them, and they are equivalent at inference time:
+
+**a) Cast on the fly** — keep the official fp32 checkpoints on disk and convert
+them while loading:
+
+```sh
+FIRERED_WEIGHT_DTYPE=bfloat16 python your_script.py
+```
+
+Nothing on disk changes (19 GB stays 19 GB), and you can go back to fp32 by
+dropping the variable. The cast is redone on every start.
+
+**b) Convert the checkpoints once** — halve them on disk too, then feed those
+weights with no variable set at all:
+
+```sh
+python scripts/convert_to_bf16.py          # in place: 19 GB -> 9.7 GB
+python your_script.py                      # loads bfloat16 by itself
+```
+
+Each tensor of the new file is verified against the fp32 original for exact
+bfloat16 rounding before the original is replaced, and the script skips
+components that are already converted. Loading gets faster too (3.3 s vs
+5.7-9.2 s here), since half as much is read from disk. Pass `--keep-fp32` to
+write `<dir>.bf16` alongside instead of replacing, and note the cast is one-way
+— recovering fp32 means re-downloading.
+
+Option (b) is also the way to move a converted model to another machine: copy
+`pretrained_models/` over, or run the script there after downloading. Either
+way, leaving `FIRERED_WEIGHT_DTYPE` unset means the loader keeps whatever the
+checkpoint stores, so nothing changes unless you ask for it.
+
+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 is off by default off CUDA.
+
 ### Model Download
 
 Download the pretrained model from Hugging Face with the `hf` CLI:

fireredtts3/campp/campp.py

@@ -2,10 +2,13 @@
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
 from fireredtts3.campp.DTDNN import CAMPPlus
+from fireredtts3.utils.device import fft_device
 
 
 def extract_kaldi_mel(audio: torch.Tensor, audio_sr: int):
     audio = audio[:1]
+    # kaldi.fbank runs an FFT internally; fall back to CPU on devices without it
+    audio = fft_device(audio)
     if audio_sr != 16000:
         audio = torchaudio.functional.resample(audio, audio_sr, 16000)
         audio_sr = 16000

fireredtts3/core.py

@@ -381,7 +381,7 @@ def generate_tts(
         gen_audio_sr = None
         segments: List[torch.Tensor] = []
         for sent in sentences:
-            seg, seg_sr, _ = super().generate_tts(
+            seg, seg_sr = super().generate_tts(
                 prompt_text=prompt_text,
                 prompt_audio=prompt_audio,
                 prompt_audio_sr=prompt_audio_sr,

fireredtts3/llm/fireredtts3_base.py

@@ -11,6 +11,7 @@
 from fireredtts3.redae.redae import RedAE
 from fireredtts3.campp.campp import CamppEmbedding
 from fireredtts3.utils.utils import fix_seed
+from fireredtts3.utils.device import get_device, get_attn_implementation, get_weight_dtype, autocast
 from fireredtts3.utils.text_tokenizer import (
     load_text_tokenizer, 
     MULTI_LANG_TAGS, MULTI_DIALECT_TAGS,
@@ -46,7 +47,7 @@
     "use_cache": True,
     "use_sliding_window": False,
     "vocab_size": 151936,
-    "attn_implementation": "flash_attention_2",
+    "attn_implementation": get_attn_implementation(),
 }
 
 
@@ -140,7 +141,7 @@ def _init_weights(self, module):
             module.rope_init()
 
     # Backbone Transformer AR wrapper
-    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
+    @autocast
     def _backbone_one_step(self, input_embeds: torch.Tensor, cache = None):
         outs = self.backbone_llm.forward(
             inputs_embeds=input_embeds,
@@ -162,7 +163,8 @@ def _flow_one_step(
         inference_cfg: float,
     ):
         # Compose input
-        x0 = torch.randn(1, self.patch_size, self.redae_dim, device=hist_latents.device)
+        x0 = torch.randn(1, self.patch_size, self.redae_dim,
+                         device=hist_latents.device, dtype=hist_latents.dtype)
 
         xt = torch.cat([hist_latents, x0], dim=1)   # History clean + current noise
         cond = torch.cat([
@@ -213,7 +215,7 @@ def generate(
         input_embeds = torch.cat([spk_embs_llm.unsqueeze(1), input_embeds, patch_prompt_latents], dim=1)
 
         # Prepare DiT decode
-        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device)
+        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device=device, dtype=input_embeds.dtype)
         t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) # (n_timesteps+1,)
         latents_gen = F.pad(prompt_latents, (0, 0, self.history_length, 0))
 
@@ -264,16 +266,16 @@ def generate(
 # RedAE + TextTokenizer + TTS3Core
 class FireRedTTS3Base(object):
     def __init__(self, pretrained_model_dir: str):
-        self.device = torch.device('cuda')
+        self.device = get_device()
         # RedAE
         redae_model_dir = os.path.join(pretrained_model_dir, 'redae')
         assert os.path.exists(redae_model_dir), f'{redae_model_dir} not found'
-        self.redae = RedAE.from_pretrained(redae_model_dir)
+        self.redae = RedAE.from_pretrained(redae_model_dir, dtype=get_weight_dtype())
         self.redae.to(self.device)
         # LLM-DiT 
         tts_model_dir = os.path.join(pretrained_model_dir, 'fireredtts3_base')
         assert os.path.exists(tts_model_dir), f'{tts_model_dir} not found'
-        self.tts_core = FireRedTTS3BaseCore.from_pretrained(tts_model_dir)
+        self.tts_core = FireRedTTS3BaseCore.from_pretrained(tts_model_dir, dtype=get_weight_dtype())
         self.tts_core.to(self.device)
         # Text Tokenizer
         text_tok_dir = os.path.join(pretrained_model_dir, 'text_tokenizer')
@@ -320,10 +322,10 @@ def generate(
         prompt_audio = self.redae.pad_to_multiple_of(prompt_audio, self.redae.downsample_rate*self.tts_core.patch_size)
         prompt_audio = prompt_audio.to(self.device)
         prompt_latents = self.redae.encode(prompt_audio, prompt_audio_sr)
-        prompt_latents = prompt_latents.to(torch.float32)
+        prompt_latents = prompt_latents.to(self.tts_core.dtype)
         # Spk emb
         spk_emb = self.spk_extractor.forward(prompt_audio, prompt_audio_sr)
-        spk_emb = spk_emb.to(self.device)
+        spk_emb = spk_emb.to(device=self.device, dtype=self.tts_core.dtype)
         # TTS
         if seed is not None:
             fix_seed(seed)

fireredtts3/llm/fireredtts3_instruct.py

@@ -17,6 +17,7 @@
 from fireredtts3.llm.dit import DiT
 from fireredtts3.redae.redae import RedAE
 from fireredtts3.utils.utils import fix_seed
+from fireredtts3.utils.device import get_device, get_weight_dtype, autocast
 from fireredtts3.utils.text_tokenizer import load_text_tokenizer
 from fireredtts3.llm.fireredtts3_base import Qwen3_1_7B_ConfigDict
 from fireredtts3.utils.chatml import (
@@ -139,7 +140,7 @@ def _init_weights(self, module):
             module.rope_init()
 
     # Backbone Transformer AR wrapper
-    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
+    @autocast
     def _backbone_one_step(self, input_embeds: torch.Tensor, cache = None):
         outs = self.backbone_llm.model.forward(
             inputs_embeds=input_embeds,
@@ -160,7 +161,8 @@ def _flow_one_step(
         inference_cfg: float,
     ):
         # Compose input
-        x0 = torch.randn(1, self.patch_size, self.redae_dim, device=hist_latents.device)
+        x0 = torch.randn(1, self.patch_size, self.redae_dim,
+                         device=hist_latents.device, dtype=hist_latents.dtype)
 
         xt = torch.cat([hist_latents, x0], dim=1)   # History clean + current noise
         cond = backbone_cond.repeat_interleave(self.patch_size, dim=1)    # Correspond backbone cond
@@ -228,11 +230,11 @@ def generate(
                 latents_patch_out.reshape(-1).to(input_embeds),
             )
         else:
-            latents_out = torch.zeros(1, 0, self.config.redae_dim, device=device)
+            latents_out = torch.zeros(1, 0, self.config.redae_dim, device=device, dtype=input_embeds.dtype)
             latents_patch_out = None
 
         # Prepare DiT decode
-        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device)
+        t_span = torch.linspace(0, 1, n_timesteps + 1).to(device=device, dtype=input_embeds.dtype)
         t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) # (n_timesteps+1,)
 
         # Init Backbone states
@@ -316,16 +318,16 @@ def generate(
 # RedAE + TextTokenizer + TTS3Core
 class FireRedTTS3Instruct(object):
     def __init__(self, pretrained_model_dir: str):
-        self.device = torch.device('cuda')
+        self.device = get_device()
         # RedAE
         redae_model_dir = os.path.join(pretrained_model_dir, 'redae')
         assert os.path.exists(redae_model_dir), f'{redae_model_dir} not found'
-        self.redae = RedAE.from_pretrained(redae_model_dir)
+        self.redae = RedAE.from_pretrained(redae_model_dir, dtype=get_weight_dtype())
         self.redae.to(self.device)
         # LLM-DiT 
         tts_model_dir = os.path.join(pretrained_model_dir, 'fireredtts3_instruct')
         assert os.path.exists(tts_model_dir), f'{tts_model_dir} not found'
-        self.tts_core = FireRedTTS3InstructCore.from_pretrained(tts_model_dir)
+        self.tts_core = FireRedTTS3InstructCore.from_pretrained(tts_model_dir, dtype=get_weight_dtype())
         self.tts_core.to(self.device)
         # Text Tokenizer
         text_tok_dir = os.path.join(pretrained_model_dir, 'text_tokenizer')
@@ -347,7 +349,7 @@ def _tokenize_audio(self, audio: torch.Tensor, audio_sr: int):
         audio = self.redae.pad_to_multiple_of(audio, self.redae.downsample_rate*self.tts_core.patch_size)
         audio = audio.to(self.device)
         latents = self.redae.encode(audio, audio_sr) * REDAE_SCALE
-        latents = latents.to(torch.float32)
+        latents = latents.to(self.tts_core.dtype)
         return latents
 
     # --- Inference Interface

fireredtts3/llm/rotary_embedding.py

@@ -1,8 +1,8 @@
 import torch
 from torch.nn import Module
-from torch.amp import autocast
 from torch import cat, stack, arange
 from einops import rearrange
+from fireredtts3.utils.device import disable_autocast
 
 
 class RotaryEmbedding(Module):
@@ -38,7 +38,7 @@ def forward_from_seq_len(self, seq_len):
         t = arange(seq_len, device = device)
         return self.forward(t)
 
-    @autocast('cuda', enabled = False)
+    @disable_autocast
     def forward(self, t, offset = 0):
         if t.ndim == 1:
             t = rearrange(t, 'n -> 1 n')
@@ -57,7 +57,7 @@ def rotate_half(x):
     return rearrange(x, '... d r -> ... (d r)')
 
 
-@autocast('cuda', enabled = False)
+@disable_autocast
 def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype
 

fireredtts3/redae/redae.py

@@ -3,6 +3,7 @@
 import torchaudio
 import torch.nn as nn
 import torch.nn.functional as F
+from fireredtts3.utils.device import get_attn_implementation, autocast, fft_device
 from transformers import (
     Qwen3Config, Qwen3Model, 
     PretrainedConfig, PreTrainedModel
@@ -34,7 +35,7 @@ def __init__(
             max_position_embeddings=max_position_embeddings,
             num_attention_heads=num_attention_heads,
             num_key_value_heads=num_key_value_heads,
-            attn_implementation='flash_attention_2',
+            attn_implementation=get_attn_implementation(),
         )
         self.qwen3 = Qwen3Model(self.qwen3_config)
         self.cls_tok = torch.nn.Parameter(torch.ones(1, 1, hidden_size))
@@ -119,7 +120,7 @@ def __init__(
             num_key_value_heads=num_key_value_heads,
             sliding_window=sliding_window,
             use_sliding_window=use_sliding_window,
-            attn_implementation='flash_attention_2',
+            attn_implementation=get_attn_implementation(),
         )
         self.qwen3 = Qwen3Model(self.qwen3_config)
         if self.extra_downsample_rate > 1:
@@ -197,9 +198,16 @@ def __init__(
         self.register_buffer("window", window)
 
     def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        # Inverse FFT / complex math is unavailable on some backends (e.g. MPS on
+        # older torch); compute on an FFT-capable device and move the result back.
+        out_device = spec.device
+        spec = fft_device(spec)
+        window = self.window.to(spec.device)
+
         if self.padding == "center":
             # Fallback to pytorch native implementation
-            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+            audio = torch.istft(spec, self.n_fft, self.hop_length, self.win_length, window, center=True)
+            return audio.to(out_device)
         elif self.padding == "same":
             pad = (self.win_length - self.hop_length) // 2
         else:
@@ -210,7 +218,7 @@ def forward(self, spec: torch.Tensor) -> torch.Tensor:
 
         # Inverse FFT
         ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
-        ifft = ifft * self.window[None, :, None]
+        ifft = ifft * window[None, :, None]
 
         # Overlap and Add
         output_size = (T - 1) * self.hop_length + self.win_length
@@ -219,7 +227,7 @@ def forward(self, spec: torch.Tensor) -> torch.Tensor:
         )[:, 0, 0, pad:-pad]
 
         # Window envelope
-        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_sq = window.square().expand(1, T, -1).transpose(1, 2)
         window_envelope = torch.nn.functional.fold(
             window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
         ).squeeze()[pad:-pad]
@@ -228,7 +236,7 @@ def forward(self, spec: torch.Tensor) -> torch.Tensor:
         assert (window_envelope > 1e-11).all()
         y = y / window_envelope
 
-        return y
+        return y.to(out_device)
 
 
 class ISTFTHead(nn.Module):
@@ -246,7 +254,8 @@ def __init__(
         self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x_pred = self.out(x)
+        # complex half is not usable -> run the spectrum/ISTFT in fp32
+        x_pred = self.out(x).float()
         x_pred = x_pred.transpose(1, 2)
         mag, p = x_pred.chunk(2, dim=1)
         mag = torch.exp(mag)
@@ -467,7 +476,7 @@ def pad_to_multiple_of(audio: torch.Tensor, multiple_of: int):
             audio = F.pad(audio, (pad_len, 0))  # NOTE left pad
         return audio
 
-    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
+    @autocast
     @torch.no_grad()
     def encode(self, audio: torch.Tensor, audio_sr:int):
         """
@@ -480,6 +489,7 @@ def encode(self, audio: torch.Tensor, audio_sr:int):
         audio = audio[:1]
         audio = torchaudio.functional.resample(audio, audio_sr, self.sample_rate)
         audio = self.pad_to_multiple_of(audio, self.downsample_rate)
+        audio = audio.to(self.dtype)
         latents = self.encoder.forward(audio)
         return latents