diff --git a/run_misotts.py b/run_misotts.py index 98b2b3b..01f93b9 100644 --- a/run_misotts.py +++ b/run_misotts.py @@ -1,25 +1,88 @@ +import argparse import os +from typing import Sequence os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60") os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60") -import torch -import torchaudio # type: ignore -from generator import DEFAULT_MISO_TTS_REPO_ID, Segment, load_miso_8b - # Disable Triton compilation os.environ["NO_TORCH_COMPILE"] = "1" -def main(): +def _positive_float(value: str) -> float: + parsed = float(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be greater than 0") + return parsed + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be greater than 0") + return parsed + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate the example Miso TTS conversation.") + parser.add_argument( + "--model", + default=os.environ.get("MISO_TTS_8B_MODEL"), + help="Local checkpoint path or Hugging Face repo id. Defaults to MISO_TTS_8B_MODEL or MisoLabs/MisoTTS.", + ) + parser.add_argument( + "--device", + default="auto", + help="Inference device, for example auto, cuda, cuda:0, or cpu. Defaults to auto.", + ) + parser.add_argument( + "--output", + default="full_conversation.wav", + help="Output WAV path. Defaults to full_conversation.wav.", + ) + parser.add_argument( + "--max-audio-length-ms", + type=_positive_float, + default=10_000, + help="Maximum generated audio length per utterance in milliseconds. Defaults to 10000.", + ) + parser.add_argument( + "--temperature", + type=_positive_float, + default=0.9, + help="Sampling temperature. Defaults to 0.9.", + ) + parser.add_argument( + "--topk", + type=_positive_int, + default=50, + help="Top-k sampling value. Defaults to 50.", + ) + return parser.parse_args(argv) + + +def resolve_device(device: str) -> str: + if device != "auto": + return device + + import torch + # Select the best available device, skipping MPS due to float64 limitations. if torch.cuda.is_available(): - device = "cuda" - else: - device = "cpu" + return "cuda" + return "cpu" + + +def main(): + args = parse_args() + device = resolve_device(args.device) print(f"Using device: {device}") - model_source = os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID) + import torch + import torchaudio # type: ignore + from generator import DEFAULT_MISO_TTS_REPO_ID, Segment, load_miso_8b + + model_source = args.model or DEFAULT_MISO_TTS_REPO_ID if os.path.exists(model_source): print(f"Loading Miso TTS model from local path: {model_source}") else: @@ -29,7 +92,7 @@ def main(): ) print("The model will be downloaded and cached automatically if it is not already present.") - generator = load_miso_8b(device, model_path_or_repo_id=model_source) + generator = load_miso_8b(device=device, model_path_or_repo_id=model_source) conversation = [ {"text": "I'm just honestly not that into him, you know?", "speaker_id": 0}, @@ -51,7 +114,9 @@ def main(): text=utterance["text"], speaker=utterance["speaker_id"], context=generated_segments, - max_audio_length_ms=10_000, + max_audio_length_ms=args.max_audio_length_ms, + temperature=args.temperature, + topk=args.topk, ) generated_segments.append( Segment( @@ -63,11 +128,11 @@ def main(): all_audio = torch.cat([seg.audio for seg in generated_segments], dim=0) torchaudio.save( - "full_conversation.wav", + args.output, all_audio.unsqueeze(0).cpu(), generator.sample_rate, ) - print("Successfully generated full_conversation.wav") + print(f"Successfully generated {args.output}") if __name__ == "__main__": diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_run_misotts.py b/tests/test_run_misotts.py new file mode 100644 index 0000000..4a6e3a4 --- /dev/null +++ b/tests/test_run_misotts.py @@ -0,0 +1,69 @@ +import io +import os +import sys +import unittest +from unittest import mock + +import run_misotts + + +class RunMisoTTSArgsTest(unittest.TestCase): + def test_defaults_match_existing_example(self): + with mock.patch.dict(os.environ, {}, clear=True): + args = run_misotts.parse_args([]) + + self.assertEqual(args.device, "auto") + self.assertEqual(args.output, "full_conversation.wav") + self.assertIsNone(args.model) + self.assertEqual(args.max_audio_length_ms, 10_000) + self.assertEqual(args.temperature, 0.9) + self.assertEqual(args.topk, 50) + + def test_model_defaults_to_environment_override(self): + with mock.patch.dict(os.environ, {"MISO_TTS_8B_MODEL": "local-checkpoint"}): + args = run_misotts.parse_args([]) + + self.assertEqual(args.model, "local-checkpoint") + + def test_accepts_custom_generation_settings(self): + args = run_misotts.parse_args( + [ + "--model", + "MisoLabs/MisoTTS", + "--device", + "cpu", + "--output", + "sample.wav", + "--max-audio-length-ms", + "5000", + "--temperature", + "0.7", + "--topk", + "25", + ] + ) + + self.assertEqual(args.model, "MisoLabs/MisoTTS") + self.assertEqual(args.device, "cpu") + self.assertEqual(args.output, "sample.wav") + self.assertEqual(args.max_audio_length_ms, 5000) + self.assertEqual(args.temperature, 0.7) + self.assertEqual(args.topk, 25) + + def test_rejects_non_positive_generation_settings(self): + with mock.patch.object(sys, "stderr", io.StringIO()): + with self.assertRaises(SystemExit): + run_misotts.parse_args(["--topk", "0"]) + + with self.assertRaises(SystemExit): + run_misotts.parse_args(["--max-audio-length-ms", "0"]) + + with self.assertRaises(SystemExit): + run_misotts.parse_args(["--temperature", "0"]) + + def test_resolve_device_returns_explicit_device(self): + self.assertEqual(run_misotts.resolve_device("cuda:1"), "cuda:1") + + +if __name__ == "__main__": + unittest.main()