Small clean up

src/python/larynx_train/export_generator.py

@@ -2,7 +2,6 @@
 import argparse
 import logging
 from pathlib import Path
-from typing import Optional
 
 import torch
 

src/python/larynx_train/infer_generator.py

@@ -8,7 +8,6 @@ from pathlib import Path
 
 import torch
 
-from .vits.lightning import VitsModel
 from .vits.utils import audio_float_to_int16
 from .vits.wavfile import write as write_wav
 
@@ -19,9 +18,7 @@ def main():
     """Main entry point"""
     logging.basicConfig(level=logging.DEBUG)
     parser = argparse.ArgumentParser(prog="larynx_train.infer_generator")
-    parser.add_argument(
-        "--model", required=True, help="Path to generator (.pt)"
-    )
+    parser.add_argument("--model", required=True, help="Path to generator (.pt)")
     parser.add_argument("--output-dir", required=True, help="Path to write WAV files")
     parser.add_argument("--sample-rate", type=int, default=22050)
     args = parser.parse_args()

src/python/larynx_train/infer_torchscript.py

@@ -8,7 +8,6 @@ from pathlib import Path
 
 import torch
 
-from .vits.lightning import VitsModel
 from .vits.utils import audio_float_to_int16
 from .vits.wavfile import write as write_wav
 

src/python/larynx_train/preprocess.py

@@ -10,7 +10,7 @@ from collections import Counter
 from dataclasses import dataclass
 from multiprocessing import JoinableQueue, Process, Queue
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional, Set
+from typing import Dict, Iterable, List, Optional
 
 from espeak_phonemizer import Phonemizer
 
@@ -20,7 +20,7 @@ from .phonemize import DEFAULT_PHONEME_ID_MAP, phonemes_to_ids, phonemize
 _LOGGER = logging.getLogger("preprocess")
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--input-dir", required=True, help="Directory with audio dataset"
@@ -124,9 +124,11 @@ def main():
     if (args.max_workers is None) or (args.max_workers < 1):
         args.max_workers = os.cpu_count()
 
+    assert args.max_workers is not None
+
     batch_size = int(num_utterances / (args.max_workers * 2))
-    queue_in = JoinableQueue()
-    queue_out = Queue()
+    queue_in: "Queue[Iterable[Utterance]]" = JoinableQueue()
+    queue_out: "Queue[Optional[Utterance]]" = Queue()
 
     # Start workers
     processes = [