220 lines
6.9 KiB
Python
220 lines
6.9 KiB
Python
import argparse
|
|
import torch.multiprocessing as mp
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
|
|
def generate_xtts_for_text(text, file, lang, speaker):
|
|
from TTS.api import TTS
|
|
import torch
|
|
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
|
|
tts.tts_to_file(
|
|
text=text,
|
|
speaker_wav=speaker,
|
|
language=lang,
|
|
file_path=file,
|
|
)
|
|
|
|
|
|
def generate_xtts(**kwargs):
|
|
import torch
|
|
import ast
|
|
import uuid
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
text = kwargs.get('text')
|
|
folder = kwargs.get('folder', '.')
|
|
folder = os.path.abspath(folder)
|
|
lang = kwargs.get('lang', 'en')
|
|
speaker = kwargs.get('speaker', '')
|
|
|
|
if isinstance(text, str):
|
|
if text != None and text.startswith("[") and text.endswith("]"):
|
|
try:
|
|
texts = ast.literal_eval(text)
|
|
except SyntaxError:
|
|
texts = [text]
|
|
elif isinstance(text, list):
|
|
texts = text
|
|
|
|
o_files = []
|
|
batch_size = 4
|
|
with ProcessPoolExecutor(max_workers=batch_size) as executor:
|
|
futures = []
|
|
for i, text in enumerate(texts):
|
|
if len(text) > 1:
|
|
ii = str(i).zfill(1 + len(str(1 + len(texts))))
|
|
file_to_use = f"{folder}/tts_{ii}_{str(uuid.uuid4())}.wav"
|
|
future = executor.submit(generate_xtts_for_text, text, file_to_use, lang, speaker)
|
|
futures.append((future, file_to_use))
|
|
|
|
for future, file_to_use in futures:
|
|
future.result()
|
|
o_files.append(file_to_use)
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
return (o_files[0], o_files,)
|
|
|
|
|
|
def cache_parameters(cache_file, params):
|
|
"""Cache the input parameters to a JSON file."""
|
|
import json
|
|
with open(cache_file, 'w') as f:
|
|
json.dump(params, f)
|
|
|
|
|
|
def load_cached_parameters(cache_file):
|
|
"""Load the cached parameters from the JSON file."""
|
|
import json
|
|
import os
|
|
if os.path.exists(cache_file):
|
|
with open(cache_file, 'r') as f:
|
|
return json.load(f)
|
|
return None
|
|
|
|
|
|
def run_test():
|
|
"""Run unit tests for the module."""
|
|
unittest.main(argv=[''], exit=False)
|
|
|
|
|
|
class TestGenerateXTTS(unittest.TestCase):
|
|
"""Unit tests for generate_xtts functions."""
|
|
|
|
@patch('os.path.exists', return_value=True)
|
|
@patch('TTS.api.TTS.tts_to_file')
|
|
def test_generate_xtts_single_text(self, mock_tts_to_file, mock_exists):
|
|
"""Test generating xtts with a single text input."""
|
|
mock_tts_to_file.return_value = None
|
|
kwargs = {
|
|
'text': 'Hello, this is a test.',
|
|
'folder': './output',
|
|
'lang': 'en',
|
|
'speaker': '/path/to/speaker.wav'
|
|
}
|
|
|
|
result = generate_xtts(**kwargs)
|
|
|
|
# Assertions
|
|
self.assertEqual(len(result[1]), 1)
|
|
mock_tts_to_file.assert_called_once()
|
|
|
|
@patch('os.path.exists', return_value=True)
|
|
@patch('TTS.api.TTS.tts_to_file')
|
|
def test_generate_xtts_multiple_texts(self, mock_tts_to_file, mock_exists):
|
|
"""Test generating xtts with multiple texts."""
|
|
mock_tts_to_file.return_value = None
|
|
kwargs = {
|
|
'text': ['Hello, this is a test.', 'This is another test.'],
|
|
'folder': './output',
|
|
'lang': 'en',
|
|
'speaker': '/path/to/speaker.wav'
|
|
}
|
|
|
|
result = generate_xtts(**kwargs)
|
|
|
|
# Assertions
|
|
self.assertEqual(len(result[1]), 2)
|
|
self.assertEqual(mock_tts_to_file.call_count, 2)
|
|
|
|
def test_cache_parameters(self):
|
|
"""Test caching parameters to a file."""
|
|
import os
|
|
import json
|
|
test_cache_file = './test_cache.json'
|
|
test_params = {'text': 'sample text', 'lang': 'en'}
|
|
|
|
cache_parameters(test_cache_file, test_params)
|
|
|
|
# Verify file content
|
|
with open(test_cache_file, 'r') as f:
|
|
data = json.load(f)
|
|
self.assertEqual(data, test_params)
|
|
|
|
# Clean up
|
|
os.remove(test_cache_file)
|
|
|
|
def test_load_cached_parameters(self):
|
|
"""Test loading cached parameters from a file."""
|
|
import os
|
|
import json
|
|
test_cache_file = './test_cache.json'
|
|
test_params = {'text': 'sample text', 'lang': 'en'}
|
|
|
|
# Write test data to the file
|
|
with open(test_cache_file, 'w') as f:
|
|
json.dump(test_params, f)
|
|
|
|
# Load the parameters
|
|
loaded_params = load_cached_parameters(test_cache_file)
|
|
self.assertEqual(loaded_params, test_params)
|
|
|
|
# Clean up
|
|
os.remove(test_cache_file)
|
|
|
|
|
|
def main():
|
|
import sys
|
|
import os
|
|
import argparse
|
|
import logging
|
|
import traceback
|
|
|
|
log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xttscli.log')
|
|
logging.basicConfig(filename=log_file, level=logging.ERROR,
|
|
format='%(asctime)s %(levelname)s %(message)s')
|
|
|
|
cache_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'params_cache.json')
|
|
|
|
try:
|
|
parser = argparse.ArgumentParser(description='Generate from given text.')
|
|
parser.add_argument('--text', help='Text to generate for')
|
|
parser.add_argument('--folder', type=str, default='.', help='Output folder')
|
|
parser.add_argument('--lang', type=str, default='en')
|
|
parser.add_argument('--speaker', type=str, default='', help='Speaker wav to use for generation')
|
|
parser.add_argument('--file', type=str, default="", help='Newline separated file.')
|
|
parser.add_argument('--rerun', action='store_true', help='Rerun with the last cached parameters.')
|
|
parser.add_argument('--test', action='store_true', help='Run the test.')
|
|
|
|
args = parser.parse_args()
|
|
kwargs = vars(args)
|
|
|
|
if args.test:
|
|
run_test()
|
|
sys.exit(0)
|
|
|
|
if args.rerun:
|
|
cached_params = load_cached_parameters(cache_file)
|
|
if cached_params:
|
|
print(f"Rerunning with cached parameters: {cached_params}")
|
|
kwargs.update(cached_params)
|
|
else:
|
|
print("No cached parameters found. Exiting.")
|
|
sys.exit(1)
|
|
else:
|
|
cache_parameters(cache_file, kwargs)
|
|
|
|
if kwargs['file']:
|
|
with open(kwargs['file']) as f:
|
|
kwargs['text'] = f.read().split('\n')
|
|
|
|
if kwargs['speaker']:
|
|
if not os.path.exists(kwargs['speaker']):
|
|
raise ValueError(f"Speaker file {kwargs['speaker']} does not exist.")
|
|
|
|
generated_files = generate_xtts(**kwargs)
|
|
print(f"Generated files: {generated_files}")
|
|
|
|
except Exception as e:
|
|
logging.error("An error occurred: %s", traceback.format_exc())
|
|
print(f"An error occurred. Check the log file at {log_file} for details.")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|