Multi-thread MP3 - Ace-Step (#2233)

* Experimental Multi-thread MP3

I used Deepseek V4 Flash to come up with a method of making the MP3 conversion with Ace-step multi-threaded.

I built and tested. Resulting MP3 doesnt have any glitches and generates almost instantly (compared to dozens of seconds previously)

* Implement WBRUNA suggestion

Updated as suggested

* further tidyup

Both issues fixed:

fprintf removed from encode_chunk (line 651) — it now just returns "" silently on failure. Error reporting moved to call sites:
Single-threaded (line 689-692): checks mp3_data.empty() and logs once
Multi-threaded (line 722-742): uses std::atomic<int> init_failures{0}, each thread does fetch_add(1) on failure, then after join a single fprintf reports the count. No interleaved output.
No leftover encode call outside threads — all chunks are encoded inside threads (lines 724-732). No results[0] = encode_chunk(...) exists between splitting and joining.

* Overlap by 1 frame

Overlapping by 1 frame from the previous chunk and discarding appears to have resolved the audible boundary glitch.

* Improve multi-threading logic in audio encoding

Refactor thread count calculation and chunk partitioning logic for audio encoding.

---------

Co-authored-by: LostRuins Concedo <39025047+LostRuins@users.noreply.github.com>

otherarch/utils.cpp

@@ -10,6 +10,9 @@
 #include <codecvt>
 #include <sstream>
 #include <ctime>
+#include <thread>
+#include <atomic>
+#include <algorithm>
 
 #define MINIAUDIO_IMPLEMENTATION
 #ifndef MTMD_AUDIO_DEBUG
@@ -635,40 +638,153 @@ std::string save_stereo_mp3_base64(const std::vector<float> & raw_audio,int T_au
         resampled = resample_wav(2,raw_audio,sample_rate,44100);
         enc_audio = resampled.data();
         enc_sr    = 44100;
+        enc_T     = (int)resampled.size() / 2;
     }
 
     const int kbps = 128;
-    mp3enc_t * enc = mp3enc_init(enc_sr, 2, kbps);
-    if (!enc) {
-        fprintf(stderr, "[Audio] mp3enc_init failed\n");
+
+    // Encode PCM [offset, offset+len) with optional warm-up.
+    // warmup=true primes filterbank/MDCT/psy with the 1152 PCM samples
+    // immediately preceding offset (caller guarantees offset >= 1152).
+    // Returns raw MP3 bytes, or empty string on init failure.
+    auto encode_chunk = [&](int offset, int len, bool warmup) -> std::string {
+        mp3enc_t * enc = mp3enc_init(enc_sr, 2, kbps);
+        if (!enc) return "";
+
+        if (warmup) {
+            int wo = offset - 1152;
+            float wb[1152 * 2];
+            memcpy(wb,          enc_audio + wo,        1152 * sizeof(float));
+            memcpy(wb + 1152,   enc_audio + enc_T + wo, 1152 * sizeof(float));
+            int dummy = 0;
+            mp3enc_encode(enc, wb, 1152, &dummy);
+            mp3enc_flush(enc, &dummy);
+        }
+
+        std::string result;
+        result.reserve(len / 4);
+
+        int chunk_size = enc_sr;
+        std::vector<float> buf((size_t)chunk_size * 2);
+
+        for (int pos = 0; pos < len; pos += chunk_size) {
+            int n = (pos + chunk_size <= len) ? chunk_size : (len - pos);
+            memcpy(buf.data(),     enc_audio + offset + pos,          (size_t)n * sizeof(float));
+            memcpy(buf.data() + n, enc_audio + enc_T + offset + pos,  (size_t)n * sizeof(float));
+            int out_size = 0;
+            const uint8_t * mp3 = mp3enc_encode(enc, buf.data(), n, &out_size);
+            if (out_size > 0) {
+                result.append((const char *) mp3, out_size);
+            }
+        }
+
+        int flush_size = 0;
+        const uint8_t * flush_data = mp3enc_flush(enc, &flush_size);
+        if (flush_size > 0) {
+            result.append((const char *) flush_data, flush_size);
+        }
+        mp3enc_free(enc);
+        return result;
+    };
+
+    // Determine thread count: at most 1 thread per ~0.5s of audio
+    int n_threads = std::max(1, (int)std::thread::hardware_concurrency());
+    int min_chunk = enc_sr / 2;
+    if (n_threads * min_chunk > enc_T) {
+        n_threads = std::max(1, enc_T / min_chunk);
+    }
+
+    // Single-threaded path — call encode_chunk directly
+    if (n_threads <= 1) {
+        std::string mp3_data = encode_chunk(0, enc_T, false);
+        if (mp3_data.empty()) {
+            fprintf(stderr, "[Audio] mp3enc_init failed\n");
+            return "";
+        }
+        return kcpp_base64_encode(mp3_data);
+    }
+
+    // Multi-threaded path: partition audio, encode chunks in parallel
+    struct Chunk { int offset; int n_samples; };
+    std::vector<Chunk> chunks(n_threads);    
+    {
+        constexpr int FRAME = 1152;
+    
+        int pos = 0;
+    
+        for (int i = 0; i < n_threads; i++) {
+            // Compute ideal boundaries
+            int start = (int)(((int64_t)enc_T * i) / n_threads);
+            int end   = (int)(((int64_t)enc_T * (i + 1)) / n_threads);
+    
+            // Align internal boundaries down to frame boundaries
+            if (i != 0) {
+                start = (start / FRAME) * FRAME;
+            }
+            if (i != n_threads - 1) {
+                end = (end / FRAME) * FRAME;
+            }
+    
+            // Enforce monotonicity
+            if (start < pos) {
+                start = pos;
+            }
+            if (end < start) {
+                end = start;
+            }
+    
+            chunks[i].offset    = start;
+            chunks[i].n_samples = end - start;
+    
+            pos = end;
+        }
+    
+        // Ensure final chunk reaches exact end
+        if (!chunks.empty()) {
+            chunks.back().n_samples = enc_T - chunks.back().offset;
+        }
+    }
+
+    std::vector<std::string> results(n_threads);
+    std::vector<std::thread> threads;
+    std::atomic<int> init_failures{0};
+
+    // RAII guard: join all threads on any exception
+    struct ThreadGuard {
+        std::vector<std::thread> & t;
+        ~ThreadGuard() { for (auto & th : t) if (th.joinable()) th.join(); }
+    } guard{threads};
+
+    for (int i = 0; i < n_threads; i++) {
+        threads.emplace_back([&, i]() {
+            if (chunks[i].n_samples <= 0) return;
+
+            // Chunk 0 starts fresh; chunks 1..N-1 get warm-up from prior chunk
+            bool warmup = (i > 0);
+            std::string r = encode_chunk(chunks[i].offset, chunks[i].n_samples, warmup);
+            if (r.empty()) {
+                init_failures.fetch_add(1, std::memory_order_relaxed);
+            } else {
+                results[i] = std::move(r);
+            }
+        });
+    }
+
+    for (auto & t : threads) {
+        t.join();
+    }
+
+    if (init_failures.load(std::memory_order_relaxed) > 0) {
+        fprintf(stderr, "[Audio] mp3enc_init failed for %d chunk(s)\n", init_failures.load());
         return "";
     }
 
     std::string mp3_data;
-    mp3_data.reserve(enc_T); // rough preallocation
+    size_t total_size = 0;
+    for (auto & r : results) total_size += r.size();
+    mp3_data.reserve(total_size);
+    for (auto & r : results) mp3_data.append(r);
 
-    int chunk = enc_sr;
-
-    // reusable buffer (replaces malloc inside loop)
-    std::vector<float> buf((size_t)chunk * 2);
-
-    for (int pos = 0; pos < enc_T; pos += chunk) {
-        int n = (pos + chunk <= enc_T) ? chunk : (enc_T - pos);
-        // build planar chunk
-        memcpy(buf.data(),     enc_audio + pos,          (size_t)n * sizeof(float));
-        memcpy(buf.data() + n, enc_audio + enc_T + pos,  (size_t)n * sizeof(float));
-        int out_size = 0;
-        const uint8_t * mp3 = mp3enc_encode(enc, buf.data(), n, &out_size);
-        if (out_size > 0) {
-            mp3_data.append((const char *) mp3, out_size);
-        }
-    }
-    int flush_size = 0;
-    const uint8_t * flush_data = mp3enc_flush(enc, &flush_size);
-    if (flush_size > 0) {
-        mp3_data.append((const char *) flush_data, flush_size);
-    }
-    mp3enc_free(enc);
     return kcpp_base64_encode(mp3_data);
 }
 
@@ -968,4 +1084,4 @@ bool FullyContainedPrefix(std::vector<int> &sequence1, std::vector<int> &sequenc
         }
     }
     return true;
-}
+}