seeing something for immiscible flow

e2_tts_pytorch/e2_tts.py

@@ -25,6 +25,8 @@
 from einops.layers.torch import Rearrange
 from einops import einsum, rearrange, repeat, reduce, pack, unpack
 
+from scipy.optimize import linear_sum_assignment
+
 from x_transformers import (
     Attention,
     FeedForward,
@@ -514,7 +516,8 @@ def __init__(
             num_gateloop_layers = 2
         ),
         mel_spec_kwargs: dict = dict(),
-        frac_lengths_mask: Tuple[float, float] = (0.7, 1.)
+        frac_lengths_mask: Tuple[float, float] = (0.7, 1.),
+        immiscible = False
     ):
         super().__init__()
 
@@ -557,6 +560,10 @@ def __init__(
         self.cond_proj_in = Linear(num_channels, dim)
         self.to_pred = Linear(dim, num_channels)
 
+        # immiscible flow - https://arxiv.org/abs/2406.12303
+
+        self.immiscible = immiscible
+
     @property
     def device(self):
         return next(self.parameters()).device
@@ -741,6 +748,13 @@ def forward(
 
         x0 = torch.randn_like(x1)
 
+        # maybe immiscible flow
+
+        if self.immiscible:
+            cost = torch.cdist(x1.flatten(1), x0.flatten(1))
+            _, reorder_indices = linear_sum_assignment(cost.cpu())
+            x0 = x0[from_numpy(reorder_indices).to(cost.device)]
+
         # t is random times from above
 
         times = torch.rand((batch,), dtype = dtype, device = self.device)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "e2-tts-pytorch"
-version = "0.3.0"
+version = "0.3.1"
 description = "E2-TTS in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -31,6 +31,7 @@ dependencies = [
     'gateloop-transformer>=0.2.2',
     'jaxtyping',
     'loguru',
+    'scipy',
     'tensorboard',
     'torch>=2.0',
     'torchdiffeq',