common : fix wav buffer detection (#1819)

diff --git a/examples/common.cpp b/examples/common.cpp
index 548156ecbc0a27ed0760e404f20220fc1e2cf988..f71c5912b9debaa06897c77cf92ca77e218ea459 100644 (file)
--- a/examples/common.cpp
+++ b/examples/common.cpp
@@ -615,6 +615,21 @@ gpt_vocab::id gpt_sample_top_k_top_p_repeat(
 
 }
 
+bool is_wav_buffer(const std::string buf) {
+    // RIFF ref: https://en.wikipedia.org/wiki/Resource_Interchange_File_Format
+    // WAV ref: https://www.mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html
+    if (buf.size() < 12 || buf.substr(0, 4) != "RIFF" || buf.substr(8, 4) != "WAVE") {
+        return false;
+    }
+
+    uint32_t chunk_size = *reinterpret_cast<const uint32_t*>(buf.data() + 4);
+    if (chunk_size + 8 != buf.size()) {
+        return false;
+    }
+
+    return true;
+}
+
 bool read_wav(const std::string & fname, std::vector<float>& pcmf32, std::vector<std::vector<float>>& pcmf32s, bool stereo) {
     drwav wav;
     std::vector<uint8_t> wav_data; // used for pipe input from stdin
@@ -639,7 +654,7 @@ bool read_wav(const std::string & fname, std::vector<float>& pcmf32, std::vector
 
         fprintf(stderr, "%s: read %zu bytes from stdin\n", __func__, wav_data.size());
     }
-    else if (fname.size() > 256 && (fname.substr(0, 4) == "RIFF" || fname.substr(8, 4) == "WAVE")) {
+    else if (is_wav_buffer(fname)) {
         if (drwav_init_memory(&wav, fname.c_str(), fname.size(), nullptr) == false) {
             fprintf(stderr, "error: failed to open WAV file from fname buffer\n");
             return false;

diff --git a/examples/common.h b/examples/common.h
index aebeb0cd4f57aa5c3177fe6bdb14c3779e366bff..09094a1b8a1c62a86fea7f4a7ee936930d4b7018 100644 (file)
--- a/examples/common.h
+++ b/examples/common.h
@@ -135,6 +135,9 @@ gpt_vocab::id gpt_sample_top_k_top_p_repeat(
 // Audio utils
 //
 
+// Check if a buffer is a WAV audio file
+bool is_wav_buffer(const std::string buf);
+
 // Read WAV audio file and store the PCM data into pcmf32
 // fname can be a buffer of WAV data instead of a filename
 // The sample rate of the audio must be equal to COMMON_SAMPLE_RATE