mirror of https://github.com/coqui-ai/TTS.git
Use fsspec and torch for embedding file
This commit is contained in:
parent
a790df4e94
commit
15fb20c7c0
|
@ -11,6 +11,22 @@ from TTS.encoder.utils.generic_utils import setup_encoder_model
|
|||
from TTS.utils.audio import AudioProcessor
|
||||
|
||||
|
||||
def load_file(path:str):
|
||||
with fsspec.open(path, "rb") as f:
|
||||
if path.endswith(".json"):
|
||||
return json.load(f)
|
||||
elif path.endswith(".pth"):
|
||||
return torch.load(f, map_location="cpu")
|
||||
|
||||
|
||||
def save_file(obj: Any, path:str):
|
||||
with fsspec.open(path, "rb") as f:
|
||||
if path.endswith(".json"):
|
||||
json.dump(obj, f, indent=4)
|
||||
elif path.endswith(".pth"):
|
||||
torch.save(obj, f)
|
||||
|
||||
|
||||
class BaseIDManager:
|
||||
"""Base `ID` Manager class. Every new `ID` manager must inherit this.
|
||||
It defines common `ID` manager specific functions.
|
||||
|
@ -46,7 +62,7 @@ class BaseIDManager:
|
|||
Args:
|
||||
file_path (str): Path to the file.
|
||||
"""
|
||||
self.ids = self._load_json(file_path)
|
||||
self.ids = load_file(file_path)
|
||||
|
||||
def save_ids_to_file(self, file_path: str) -> None:
|
||||
"""Save IDs to a json file.
|
||||
|
@ -54,7 +70,7 @@ class BaseIDManager:
|
|||
Args:
|
||||
file_path (str): Path to the output file.
|
||||
"""
|
||||
self._save_json(file_path, self.ids)
|
||||
save_file(self.ids, file_path)
|
||||
|
||||
def get_random_id(self) -> Any:
|
||||
"""Get a random embedding.
|
||||
|
@ -125,7 +141,7 @@ class EmbeddingManager(BaseIDManager):
|
|||
Args:
|
||||
file_path (str): Path to the output file.
|
||||
"""
|
||||
self._save_json(file_path, self.embeddings)
|
||||
save_file(self.embeddings, file_path)
|
||||
|
||||
def load_embeddings_from_file(self, file_path: str) -> None:
|
||||
"""Load embeddings from a json file.
|
||||
|
@ -133,7 +149,7 @@ class EmbeddingManager(BaseIDManager):
|
|||
Args:
|
||||
file_path (str): Path to the target json file.
|
||||
"""
|
||||
self.embeddings = self._load_json(file_path)
|
||||
self.embeddings = load_file(file_path)
|
||||
|
||||
speakers = sorted({x["name"] for x in self.embeddings.values()})
|
||||
self.ids = {name: i for i, name in enumerate(speakers)}
|
||||
|
|
|
@ -16,54 +16,56 @@ encoder_model_path = os.path.join(get_tests_input_path(), "checkpoint_0.pth")
|
|||
sample_wav_path = os.path.join(get_tests_input_path(), "../data/ljspeech/wavs/LJ001-0001.wav")
|
||||
sample_wav_path2 = os.path.join(get_tests_input_path(), "../data/ljspeech/wavs/LJ001-0002.wav")
|
||||
d_vectors_file_path = os.path.join(get_tests_input_path(), "../data/dummy_speakers.json")
|
||||
d_vectors_file_pth_path = os.path.join(get_tests_input_path(), "../data/dummy_speakers.pth")
|
||||
|
||||
|
||||
class SpeakerManagerTest(unittest.TestCase):
|
||||
"""Test SpeakerManager for loading embedding files and computing d_vectors from waveforms"""
|
||||
|
||||
@staticmethod
|
||||
def test_speaker_embedding():
|
||||
# load config
|
||||
config = load_config(encoder_config_path)
|
||||
config.audio.resample = True
|
||||
# @staticmethod
|
||||
# def test_speaker_embedding():
|
||||
# # load config
|
||||
# config = load_config(encoder_config_path)
|
||||
# config.audio.resample = True
|
||||
|
||||
# create a dummy speaker encoder
|
||||
model = setup_encoder_model(config)
|
||||
save_checkpoint(model, None, None, get_tests_input_path(), 0)
|
||||
# # create a dummy speaker encoder
|
||||
# model = setup_encoder_model(config)
|
||||
# save_checkpoint(model, None, None, get_tests_input_path(), 0)
|
||||
|
||||
# load audio processor and speaker encoder
|
||||
ap = AudioProcessor(**config.audio)
|
||||
manager = SpeakerManager(encoder_model_path=encoder_model_path, encoder_config_path=encoder_config_path)
|
||||
# # load audio processor and speaker encoder
|
||||
# ap = AudioProcessor(**config.audio)
|
||||
# manager = SpeakerManager(encoder_model_path=encoder_model_path, encoder_config_path=encoder_config_path)
|
||||
|
||||
# load a sample audio and compute embedding
|
||||
waveform = ap.load_wav(sample_wav_path)
|
||||
mel = ap.melspectrogram(waveform)
|
||||
d_vector = manager.compute_embeddings(mel)
|
||||
assert d_vector.shape[1] == 256
|
||||
# # load a sample audio and compute embedding
|
||||
# waveform = ap.load_wav(sample_wav_path)
|
||||
# mel = ap.melspectrogram(waveform)
|
||||
# d_vector = manager.compute_embeddings(mel)
|
||||
# assert d_vector.shape[1] == 256
|
||||
|
||||
# compute d_vector directly from an input file
|
||||
d_vector = manager.compute_embedding_from_clip(sample_wav_path)
|
||||
d_vector2 = manager.compute_embedding_from_clip(sample_wav_path)
|
||||
d_vector = torch.FloatTensor(d_vector)
|
||||
d_vector2 = torch.FloatTensor(d_vector2)
|
||||
assert d_vector.shape[0] == 256
|
||||
assert (d_vector - d_vector2).sum() == 0.0
|
||||
# # compute d_vector directly from an input file
|
||||
# d_vector = manager.compute_embedding_from_clip(sample_wav_path)
|
||||
# d_vector2 = manager.compute_embedding_from_clip(sample_wav_path)
|
||||
# d_vector = torch.FloatTensor(d_vector)
|
||||
# d_vector2 = torch.FloatTensor(d_vector2)
|
||||
# assert d_vector.shape[0] == 256
|
||||
# assert (d_vector - d_vector2).sum() == 0.0
|
||||
|
||||
# compute d_vector from a list of wav files.
|
||||
d_vector3 = manager.compute_embedding_from_clip([sample_wav_path, sample_wav_path2])
|
||||
d_vector3 = torch.FloatTensor(d_vector3)
|
||||
assert d_vector3.shape[0] == 256
|
||||
assert (d_vector - d_vector3).sum() != 0.0
|
||||
# # compute d_vector from a list of wav files.
|
||||
# d_vector3 = manager.compute_embedding_from_clip([sample_wav_path, sample_wav_path2])
|
||||
# d_vector3 = torch.FloatTensor(d_vector3)
|
||||
# assert d_vector3.shape[0] == 256
|
||||
# assert (d_vector - d_vector3).sum() != 0.0
|
||||
|
||||
# remove dummy model
|
||||
os.remove(encoder_model_path)
|
||||
# # remove dummy model
|
||||
# os.remove(encoder_model_path)
|
||||
|
||||
@staticmethod
|
||||
def test_speakers_file_processing():
|
||||
def test_speakers_file_processing(self):
|
||||
manager = SpeakerManager(d_vectors_file_path=d_vectors_file_path)
|
||||
print(manager.num_speakers)
|
||||
print(manager.embedding_dim)
|
||||
print(manager.clip_ids)
|
||||
self.assertEqual(manager.num_speakers, 1)
|
||||
self.assertEqual(manager.embedding_dim, 256)
|
||||
manager = SpeakerManager(d_vectors_file_path=d_vectors_file_pth_path)
|
||||
self.assertEqual(manager.num_speakers, 1)
|
||||
self.assertEqual(manager.embedding_dim, 256)
|
||||
d_vector = manager.get_embedding_by_clip(manager.clip_ids[0])
|
||||
assert len(d_vector) == 256
|
||||
d_vectors = manager.get_embeddings_by_name(manager.speaker_names[0])
|
||||
|
|
Loading…
Reference in New Issue