Fix long text generation.

f5_tts_mlx/duration.py

@@ -183,6 +183,10 @@ def __call__(
                 text = list_str_to_tensor(text)
             assert text.shape[0] == batch
 
+        if seq_len < text.shape[1]:
+            seq_len = text.shape[1]
+            inp = mx.pad(inp, [(0, 0), (0, seq_len - inp.shape[1]), (0, 0)])
+
         # lens and mask
         if not exists(lens):
             lens = mx.full((batch,), seq_len)

f5_tts_mlx/utils.py

@@ -89,18 +89,14 @@ def maybe_masked_mean(t: mx.array, mask: mx.array | None = None) -> mx.array:
     return einx.divide("b d, b -> b d", num, mx.maximum(den, 1))
 
 
-def pad_to_length(t: mx.array, length: int, value=None):
+def pad_to_length(t: mx.array, length: int, value=0):
     ndim = t.ndim
     seq_len = t.shape[-1]
     if length > seq_len:
         if ndim == 1:
             t = mx.pad(t, [(0, length - seq_len)], constant_values=value)
         elif ndim == 2:
             t = mx.pad(t, [(0, 0), (0, length - seq_len)], constant_values=value)
-        elif ndim == 3:
-            t = mx.pad(
-                t, [(0, 0), (0, length - seq_len), (0, 0)], constant_values=value
-            )
         else:
             raise ValueError(f"Unsupported padding dims: {ndim}")
     return t[..., :length]

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "f5-tts-mlx"
-version = "0.1.5"
+version = "0.1.6"
 authors = [{name = "Lucas Newman", email = "[email protected]"}]
 license = {text = "MIT"}
 description = "F5-TTS - MLX"