it does not work

e2_tts_pytorch/e2_tts.py

@@ -63,7 +63,7 @@ def __getitem__(self, shapes: str):
 
 # named tuples
 
-LossBreakdown = namedtuple('LossBreakdown', ['flow', 'velocity_consistency', 'direction'])
+LossBreakdown = namedtuple('LossBreakdown', ['flow', 'velocity_consistency'])
 
 E2TTSReturn = namedtuple('E2TTS', ['loss', 'cond', 'pred_flow', 'pred_data', 'loss_breakdown'])
 
@@ -109,12 +109,6 @@ def project(x, y):
 
     return inverse(parallel).to(dtype), inverse(orthogonal).to(dtype)
 
-# losses
-
-def calc_direction_loss(pred, target):
-     # make direction loss at most 1.
-    return 0.5 * (1. - einsum(l2norm(pred), l2norm(target), '... d, ... d -> ...'))
-
 # simple utf-8 tokenizer, since paper went character based
 
 def list_str_to_tensor(
@@ -918,9 +912,7 @@ def __init__(
         use_vocos = True,
         pretrained_vocos_path = 'charactr/vocos-mel-24khz',
         sampling_rate: int | None = None,
-        add_direction_loss = False,
         velocity_consistency_weight = 0.,
-        direction_loss_weight = 1.
     ):
         super().__init__()
 
@@ -997,11 +989,6 @@ def __init__(
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
         self.velocity_consistency_weight = velocity_consistency_weight
 
-        # direction loss for flow matching
-
-        self.add_direction_loss = add_direction_loss
-        self.direction_loss_weight = direction_loss_weight
-
         # default vocos for mel -> audio
 
         self.vocos = Vocos.from_pretrained(pretrained_vocos_path) if use_vocos else None
@@ -1333,24 +1320,14 @@ def forward(
 
         loss = loss[rand_span_mask].mean()
 
-        # maybe direction loss
-
-        direction_loss = self.zero
-
-        if self.add_direction_loss:
-            direction_loss = calc_direction_loss(pred, flow)
-
-            direction_loss = direction_loss[rand_span_mask].mean()
-
         # total loss and get breakdown
 
         total_loss = (
             loss +
-            direction_loss * self.direction_loss_weight +
             velocity_loss * self.velocity_consistency_weight
         )
 
-        breakdown = LossBreakdown(loss, velocity_loss, direction_loss)
+        breakdown = LossBreakdown(loss, velocity_loss)
 
         # return total loss and bunch of intermediates
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "e2-tts-pytorch"
-version = "1.5.1"
+version = "1.5.2"
 description = "E2-TTS in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }