diff --git a/modules/xttscli.py b/modules/xttscli.py index f666fd1..75771ac 100644 --- a/modules/xttscli.py +++ b/modules/xttscli.py @@ -1,7 +1,8 @@ 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): @@ -13,7 +14,6 @@ def generate_xtts_for_text(text, file, lang, speaker): speaker_wav=speaker, language=lang, file_path=file, - # emotion="angry", # doesnt do anything as far as i can tell with xtts2 ) @@ -28,7 +28,6 @@ def generate_xtts(**kwargs): text = kwargs.get('text') folder = kwargs.get('folder', '.') - # make the folder absolute folder = os.path.abspath(folder) lang = kwargs.get('lang', 'en') speaker = kwargs.get('speaker', '') @@ -62,34 +61,103 @@ def generate_xtts(**kwargs): return (o_files[0], o_files,) -# def generate_xtts_for_text_advanced(text, file, lang, speaker): -# from TTS.api import TTS -# import torch -# -# # Create a SharedMemoryManager -# manager = mp.Manager() -# # Create a shared list to store the result -# result = manager.list() -# model = manager.model() -# -# -# # also we want to SHARE the TTS model -# something = manager. - 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 @@ -97,12 +165,10 @@ def main(): import logging import traceback - # Set up logging to a file 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') - # Define the cache file path cache_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'params_cache.json') try: @@ -113,11 +179,15 @@ def main(): 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) - # Handle --rerun logic + if args.test: + run_test() + sys.exit(0) + if args.rerun: cached_params = load_cached_parameters(cache_file) if cached_params: @@ -127,25 +197,20 @@ def main(): print("No cached parameters found. Exiting.") sys.exit(1) else: - # Cache the parameters if not rerun cache_parameters(cache_file, kwargs) - # If file is set, open it and read the content if kwargs['file']: with open(kwargs['file']) as f: kwargs['text'] = f.read().split('\n') - # Ensure the speaker file exists if kwargs['speaker']: if not os.path.exists(kwargs['speaker']): raise ValueError(f"Speaker file {kwargs['speaker']} does not exist.") - # Call your generate function generated_files = generate_xtts(**kwargs) print(f"Generated files: {generated_files}") except Exception as e: - # Log the full exception details including the stack trace logging.error("An error occurred: %s", traceback.format_exc()) print(f"An error occurred. Check the log file at {log_file} for details.")