beartype for init

e2_tts_pytorch/e2_tts.py

@@ -14,6 +14,7 @@
 from itertools import zip_longest
 from collections import namedtuple
 from typing import Literal, Callable
+from beartype import beartype
 
 import torch
 from torch import nn, tensor, from_numpy
@@ -304,6 +305,7 @@ def forward(
 # for use in both e2tts as well as duration module
 
 class Transformer(Module):
+    @beartype
     def __init__(
         self,
         *,
@@ -569,14 +571,18 @@ def forward(
 # main classes
 
 class DurationPredictor(Module):
+    @beartype
     def __init__(
         self,
         transformer: dict | Transformer,
         num_channels = None,
         mel_spec_kwargs: dict = dict(),
         char_embed_kwargs: dict = dict(),
         text_num_embeds = None,
-        tokenizer: str |  Callable[[list[str]], Int['b nt']] = 'char_utf8'
+        tokenizer: (
+            Literal['char_utf8', 'phoneme_en'] |
+            Callable[[list[str]], Int['b nt']]
+        ) = 'char_utf8'
     ):
         super().__init__()
 
@@ -689,6 +695,8 @@ def forward(
         return F.mse_loss(pred, lens.float())
 
 class E2TTS(Module):
+
+    @beartype
     def __init__(
         self,
         transformer: dict | Transformer = None,
@@ -706,7 +714,10 @@ def __init__(
         frac_lengths_mask: tuple[float, float] = (0.7, 1.),
         concat_cond = False,
         text_num_embeds = None,
-        tokenizer: str |  Callable[[list[str]], Int['b nt']] = 'char_utf8'
+        tokenizer: (
+            Literal['char_utf8', 'phoneme_en'] |
+            Callable[[list[str]], Int['b nt']]
+        ) = 'char_utf8'
     ):
         super().__init__()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "e2-tts-pytorch"
-version = "0.9.3"
+version = "0.9.4"
 description = "E2-TTS in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,6 +25,7 @@ classifiers=[
 
 dependencies = [
     'accelerate>=0.33.0',
+    'beartype',
     'einops>=0.8.0',
     'einx>=0.3.0',
     'ema-pytorch>=0.5.2',