add text chunking for long-form generation (Fixes #4)

swift/README.md

@@ -2,6 +2,10 @@
 
 This guide provides examples for running TTS inference using `example_onnx`.
 
+## 📰 Update News
+
+**2025.11.19** - Added automatic text chunking for long-form inference. Long texts are split into chunks and synthesized with natural pauses.
+
 ## Installation
 
 This project uses Swift Package Manager (SPM) for dependency management.
@@ -34,6 +38,7 @@ This will use:
 Process multiple voice styles and texts at once:
 ```bash
 .build/release/example_onnx \
+  --batch \
   --voice-style assets/voice_styles/M1.json,assets/voice_styles/F1.json \
   --text "The sun sets behind the mountains, painting the sky in shades of pink and orange.|The weather is beautiful and sunny outside. A gentle breeze makes the air feel fresh and pleasant."
 ```
@@ -57,6 +62,23 @@ This will:
 - Use 10 denoising steps instead of the default 5
 - Produce higher quality output at the cost of slower inference
 
+### Example 4: Long-Form Inference
+The system automatically chunks long texts into manageable segments, synthesizes each segment separately, and concatenates them with natural pauses (0.3 seconds by default) into a single audio file. This happens by default when you don't use the `--batch` flag:
+
+```bash
+.build/release/example_onnx \
+  --voice-style assets/voice_styles/M1.json \
+  --text "This is a very long text that will be automatically split into multiple chunks. The system will process each chunk separately and then concatenate them together with natural pauses between segments. This ensures that even very long texts can be processed efficiently while maintaining natural speech flow and avoiding memory issues."
+```
+
+This will:
+- Automatically split the text into chunks based on paragraph and sentence boundaries
+- Synthesize each chunk separately
+- Add 0.3 seconds of silence between chunks for natural pauses
+- Concatenate all chunks into a single audio file
+
+**Note**: Automatic text chunking is disabled when using `--batch` mode. In batch mode, each text is processed as-is without chunking.
+
 ## Available Arguments
 
 | Argument | Type | Default | Description |
@@ -68,9 +90,11 @@ This will:
 | `--voice-style` | str+ | `assets/voice_styles/M1.json` | Voice style file path(s) |
 | `--text` | str+ | (long default text) | Text(s) to synthesize |
 | `--save-dir` | str | `results` | Output directory |
+| `--batch` | flag | False | Enable batch mode (multiple text-style pairs, disables automatic chunking) |
 
 ## Notes
 
-- **Batch Processing**: The number of `--voice-style` files must match the number of `--text` entries
+- **Batch Processing**: When using `--batch`, the number of `--voice-style` files must match the number of `--text` entries
+- **Automatic Chunking**: Without `--batch`, long texts are automatically split and concatenated with 0.3s pauses
 - **Quality vs Speed**: Higher `--total-step` values produce better quality but take longer
 - **GPU Support**: GPU mode is not supported yet

swift/Sources/ExampleONNX.swift

@@ -9,6 +9,7 @@ struct Args {
     var voiceStyle: [String] = ["assets/voice_styles/M1.json"]
     var text: [String] = ["This morning, I took a walk in the park, and the sound of the birds and the breeze was so pleasant that I stopped for a long time just to listen."]
     var saveDir: String = "results"
+    var batch: Bool = false
 }
 
 func parseArgs() -> Args {
@@ -52,6 +53,8 @@ func parseArgs() -> Args {
                 args.saveDir = arguments[i + 1]
                 i += 1
             }
+        case "--batch":
+            args.batch = true
         default:
             break
         }
@@ -70,9 +73,11 @@ struct ExampleONNX {
         // --- 1. Parse arguments --- //
         let args = parseArgs()
         
-        guard args.voiceStyle.count == args.text.count else {
-            print("Error: Number of voice styles (\(args.voiceStyle.count)) must match number of texts (\(args.text.count))")
-            return
+        if args.batch {
+            guard args.voiceStyle.count == args.text.count else {
+                print("Error: Number of voice styles (\(args.voiceStyle.count)) must match number of texts (\(args.text.count))")
+                return
+            }
         }
         
         let bsz = args.voiceStyle.count
@@ -92,19 +97,39 @@ struct ExampleONNX {
             for n in 0..<args.nTest {
                 print("\n[\(n + 1)/\(args.nTest)] Starting synthesis...")
                 
-                let (wav, duration) = try timer("Generating speech from text") {
-                    try textToSpeech.call(args.text, style, args.totalStep)
+                let wav: [Float]
+                let duration: [Float]
+                
+                if args.batch {
+                    let result = try timer("Generating speech from text") {
+                        try textToSpeech.batch(args.text, style, args.totalStep)
+                    }
+                    wav = result.wav
+                    duration = result.duration
+                } else {
+                    let result = try timer("Generating speech from text") {
+                        try textToSpeech.call(args.text[0], style, args.totalStep, silenceDuration: 0.3)
+                    }
+                    wav = result.wav
+                    duration = [result.duration]
                 }
                 
                 // Save outputs
-                let wavLen = wav.count / bsz
                 for i in 0..<bsz {
                     let fname = "\(sanitizeFilename(args.text[i], maxLen: 20))_\(n + 1).wav"
-                    let actualLen = Int(Float(textToSpeech.sampleRate) * duration[i])
+                    let wavOut: [Float]
                     
-                    let wavStart = i * wavLen
-                    let wavEnd = min(wavStart + actualLen, wavStart + wavLen)
-                    let wavOut = Array(wav[wavStart..<wavEnd])
+                    if args.batch {
+                        let wavLen = wav.count / bsz
+                        let actualLen = Int(Float(textToSpeech.sampleRate) * duration[i])
+                        let wavStart = i * wavLen
+                        let wavEnd = min(wavStart + actualLen, wavStart + wavLen)
+                        wavOut = Array(wav[wavStart..<wavEnd])
+                    } else {
+                        // For non-batch mode, wav is a single concatenated audio
+                        let actualLen = Int(Float(textToSpeech.sampleRate) * duration[0])
+                        wavOut = Array(wav.prefix(actualLen))
+                    }
                     
                     let outputPath = "\(args.saveDir)/\(fname)"
                     try writeWavFile(outputPath, wavOut, textToSpeech.sampleRate)

swift/Sources/Helper.swift

@@ -203,6 +203,199 @@ func writeWavFile(_ filename: String, _ audioData: [Float], _ sampleRate: Int) t
     try data.write(to: url)
 }
 
+// MARK: - Text Chunking
+
+let MAX_CHUNK_LENGTH = 300
+let ABBREVIATIONS = [
+    "Dr.", "Mr.", "Mrs.", "Ms.", "Prof.", "Sr.", "Jr.",
+    "St.", "Ave.", "Rd.", "Blvd.", "Dept.", "Inc.", "Ltd.",
+    "Co.", "Corp.", "etc.", "vs.", "i.e.", "e.g.", "Ph.D."
+]
+
+func chunkText(_ text: String, maxLen: Int = 0) -> [String] {
+    let actualMaxLen = maxLen > 0 ? maxLen : MAX_CHUNK_LENGTH
+    let trimmedText = text.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
+    
+    if trimmedText.isEmpty {
+        return [""]
+    }
+    
+    // Split by paragraphs using regex
+    let paraPattern = try! NSRegularExpression(pattern: "\\n\\s*\\n")
+    let paraRange = NSRange(trimmedText.startIndex..., in: trimmedText)
+    var paragraphs = [String]()
+    var lastEnd = trimmedText.startIndex
+    
+    paraPattern.enumerateMatches(in: trimmedText, range: paraRange) { match, _, _ in
+        if let match = match, let range = Range(match.range, in: trimmedText) {
+            paragraphs.append(String(trimmedText[lastEnd..<range.lowerBound]))
+            lastEnd = range.upperBound
+        }
+    }
+    if lastEnd < trimmedText.endIndex {
+        paragraphs.append(String(trimmedText[lastEnd...]))
+    }
+    if paragraphs.isEmpty {
+        paragraphs = [trimmedText]
+    }
+    
+    var chunks = [String]()
+    
+    for para in paragraphs {
+        let trimmedPara = para.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
+        if trimmedPara.isEmpty {
+            continue
+        }
+        
+        if trimmedPara.count <= actualMaxLen {
+            chunks.append(trimmedPara)
+            continue
+        }
+        
+        // Split by sentences
+        let sentences = splitSentences(trimmedPara)
+        var current = ""
+        var currentLen = 0
+        
+        for sentence in sentences {
+            let trimmedSentence = sentence.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
+            if trimmedSentence.isEmpty {
+                continue
+            }
+            
+            let sentenceLen = trimmedSentence.count
+            if sentenceLen > actualMaxLen {
+                // If sentence is longer than maxLen, split by comma or space
+                if !current.isEmpty {
+                    chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+                    current = ""
+                    currentLen = 0
+                }
+                
+                // Try splitting by comma
+                let parts = trimmedSentence.components(separatedBy: ",")
+                for part in parts {
+                    let trimmedPart = part.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
+                    if trimmedPart.isEmpty {
+                        continue
+                    }
+                    
+                    let partLen = trimmedPart.count
+                    if partLen > actualMaxLen {
+                        // Split by space as last resort
+                        let words = trimmedPart.components(separatedBy: CharacterSet.whitespaces).filter { !$0.isEmpty }
+                        var wordChunk = ""
+                        var wordChunkLen = 0
+                        
+                        for word in words {
+                            let wordLen = word.count
+                            if wordChunkLen + wordLen + 1 > actualMaxLen && !wordChunk.isEmpty {
+                                chunks.append(wordChunk.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+                                wordChunk = ""
+                                wordChunkLen = 0
+                            }
+                            
+                            if !wordChunk.isEmpty {
+                                wordChunk += " "
+                                wordChunkLen += 1
+                            }
+                            wordChunk += word
+                            wordChunkLen += wordLen
+                        }
+                        
+                        if !wordChunk.isEmpty {
+                            chunks.append(wordChunk.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+                        }
+                    } else {
+                        if currentLen + partLen + 1 > actualMaxLen && !current.isEmpty {
+                            chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+                            current = ""
+                            currentLen = 0
+                        }
+                        
+                        if !current.isEmpty {
+                            current += ", "
+                            currentLen += 2
+                        }
+                        current += trimmedPart
+                        currentLen += partLen
+                    }
+                }
+                continue
+            }
+            
+            if currentLen + sentenceLen + 1 > actualMaxLen && !current.isEmpty {
+                chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+                current = ""
+                currentLen = 0
+            }
+            
+            if !current.isEmpty {
+                current += " "
+                currentLen += 1
+            }
+            current += trimmedSentence
+            currentLen += sentenceLen
+        }
+        
+        if !current.isEmpty {
+            chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
+        }
+    }
+    
+    return chunks.isEmpty ? [""] : chunks
+}
+
+func splitSentences(_ text: String) -> [String] {
+    // Swift's regex doesn't support lookbehind reliably, so we use a simpler approach
+    // Split on sentence boundaries and then check if they're abbreviations
+    let regex = try! NSRegularExpression(pattern: "([.!?])\\s+")
+    let range = NSRange(text.startIndex..., in: text)
+    
+    // Find all matches
+    let matches = regex.matches(in: text, range: range)
+    if matches.isEmpty {
+        return [text]
+    }
+    
+    var sentences = [String]()
+    var lastEnd = text.startIndex
+    
+    for match in matches {
+        guard let matchRange = Range(match.range, in: text) else { continue }
+        
+        // Get the text before the punctuation
+        let beforePunc = String(text[lastEnd..<matchRange.lowerBound])
+        
+        // Get the punctuation character
+        let puncRange = Range(NSRange(location: match.range.location, length: 1), in: text)!
+        let punc = String(text[puncRange])
+        
+        // Check if this ends with an abbreviation
+        var isAbbrev = false
+        let combined = beforePunc.trimmingCharacters(in: CharacterSet.whitespaces) + punc
+        for abbrev in ABBREVIATIONS {
+            if combined.hasSuffix(abbrev) {
+                isAbbrev = true
+                break
+            }
+        }
+        
+        if !isAbbrev {
+            // This is a real sentence boundary
+            sentences.append(String(text[lastEnd..<matchRange.upperBound]))
+            lastEnd = matchRange.upperBound
+        }
+    }
+    
+    // Add the remaining text
+    if lastEnd < text.endIndex {
+        sentences.append(String(text[lastEnd...]))
+    }
+    
+    return sentences.isEmpty ? [text] : sentences
+}
+
 // MARK: - Utility Functions
 
 func timer<T>(_ name: String, _ f: () throws -> T) rethrows -> T {
@@ -260,7 +453,7 @@ class TextToSpeech {
         self.sampleRate = cfgs.ae.sample_rate
     }
     
-    func call(_ textList: [String], _ style: Style, _ totalStep: Int) throws -> (wav: [Float], duration: [Float]) {
+    private func _infer(_ textList: [String], _ style: Style, _ totalStep: Int) throws -> (wav: [Float], duration: [Float]) {
         let bsz = textList.count
         
         // Process text
@@ -382,6 +575,39 @@ class TextToSpeech {
         
         return (wav, duration)
     }
+    
+    func call(_ text: String, _ style: Style, _ totalStep: Int, silenceDuration: Float) throws -> (wav: [Float], duration: Float) {
+        let chunks = chunkText(text)
+        
+        var wavCat = [Float]()
+        var durCat: Float = 0.0
+        
+        for (i, chunk) in chunks.enumerated() {
+            let result = try _infer([chunk], style, totalStep)
+            
+            let dur = result.duration[0]
+            let wavLen = Int(Float(sampleRate) * dur)
+            let wavChunk = Array(result.wav.prefix(wavLen))
+            
+            if i == 0 {
+                wavCat = wavChunk
+                durCat = dur
+            } else {
+                let silenceLen = Int(silenceDuration * Float(sampleRate))
+                let silence = [Float](repeating: 0.0, count: silenceLen)
+                
+                wavCat.append(contentsOf: silence)
+                wavCat.append(contentsOf: wavChunk)
+                durCat += silenceDuration + dur
+            }
+        }
+        
+        return (wavCat, durCat)
+    }
+    
+    func batch(_ textList: [String], _ style: Style, _ totalStep: Int) throws -> (wav: [Float], duration: [Float]) {
+        return try _infer(textList, style, totalStep)
+    }
 }
 
 // MARK: - Component Loading Functions