Merge pull request #17 from kahrendt/mps-support-v3

Add support for MPS (Apple Silicon) backend
This commit is contained in:
Michael Hansen 2025-09-19 15:09:08 -05:00 committed by GitHub
commit 395aef7dbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,5 +1,6 @@
#!/usr/bin/env python3
import argparse
import gc
import itertools as it
import json
import logging
@ -73,6 +74,10 @@ def generate_samples(
if torch.cuda.is_available():
torch_model.cuda()
_LOGGER.debug("CUDA available, using GPU")
elif torch.backends.mps.is_available():
mps_device = torch.device("mps")
torch_model.to(mps_device)
_LOGGER.debug("MPS available, using GPU")
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
@ -162,6 +167,11 @@ def generate_samples(
.numpy()
)
if torch.backends.mps.is_available():
# There seems to be a memory leak if we don't empty the cache after each batch with mps
torch.mps.empty_cache()
gc.collect()
audio_int16 = audio_float_to_int16(audio)
for audio_idx in range(audio_int16.shape[0]):
audio_data = np.trim_zeros(audio_int16[audio_idx].flatten())
@ -319,6 +329,12 @@ def generate_audio(
speaker_2 = speaker_2.cuda()
x = cast(torch.LongTensor, x.cuda())
x_lengths = cast(torch.LongTensor, x_lengths.cuda())
elif torch.backends.mps.is_available():
mps_device = torch.device("mps")
speaker_1 = speaker_1.to(mps_device)
speaker_2 = speaker_2.to(mps_device)
x = x.to(mps_device)
x_lengths = x_lengths.to(mps_device)
x, m_p_orig, logs_p_orig, x_mask = model.enc_p(x, x_lengths)
emb0 = model.emb_g(speaker_1)