Eval after each ODE step for accurate progress.

Author: lucasnewman

Sources/F5TTS/DiT.swift

@@ -77,9 +77,7 @@ class InputEmbedding: Module {
 
         let combined = MLX.concatenated([x, cond, textEmbed], axis: -1)
         var output = proj(combined)
-        output.eval()
         output = conv_pos_embed(output) + output
-        output.eval()
         return output
     }
 }

Sources/F5TTS/F5TTS.swift

@@ -141,7 +141,10 @@ public class F5TTS: Module {
 
             progressHandler?(Double(t))
 
-            return pred + (pred - nullPred) * cfgStrength
+            let output = pred + (pred - nullPred) * cfgStrength
+            output.eval()
+            
+            return output
         }
 
         // noise input
@@ -165,11 +168,10 @@ public class F5TTS: Module {
         let trajectory = self.odeint(fun: fn, y0: y0Padded, t: t)
         let sampled = trajectory[-1]
         var out = MLX.where(condMask, cond, sampled)
-
+        
         if let vocoder = vocoder {
             out = vocoder(out)
         }
-
         out.eval()
 
         return (out, trajectory)
@@ -239,6 +241,8 @@ public class F5TTS: Module {
         }
 
         let generatedAudio = outputAudio[audio.shape[0]...]
+        
+        print("Got generated audio of shape: \(generatedAudio.shape)")
         return generatedAudio
     }
 }
@@ -331,9 +335,8 @@ public extension F5TTS {
 
     static func estimatedDuration(refAudio: MLXArray, refText: String, text: String, speed: Double = 1.0) -> TimeInterval {
         let refDurationInFrames = refAudio.shape[0] / self.hopLength
-        let pausePunctuation = "。，、；：？！"
-        let refTextLength = refText.utf8.count + 3 * pausePunctuation.utf8.count
-        let genTextLength = text.utf8.count + 3 * pausePunctuation.utf8.count
+        let refTextLength = refText.utf8.count
+        let genTextLength = text.utf8.count
 
         let refAudioToTextRatio = Double(refDurationInFrames) / Double(refTextLength)
         let textLength = Double(genTextLength) / speed

Sources/F5TTS/Modules.swift

@@ -58,10 +58,7 @@ func precomputeFreqsCis(dim: Int, end: Int, theta: Float = 10000.0, thetaRescale
     let freqsCos = outerFreqs.cos()
     let freqsSin = outerFreqs.sin()
 
-    let output = MLX.concatenated([freqsCos, freqsSin], axis: -1)
-    output.eval()
-
-    return output
+    return MLX.concatenated([freqsCos, freqsSin], axis: -1)
 }
 
 func getPosEmbedIndices(start: MLXArray, length: Int, maxPos: Int, scale: Float = 1.0) -> MLXArray {