working converter

Makefile

@@ -8,10 +8,10 @@ PRIV_DIR = $(MIX_APP_PATH)/priv
 XAV_SO = $(PRIV_DIR)/libxav.so
 
 # uncomment to compile with debug logs
-# XAV_DEBUG_LOGS = -DXAV_DEBUG=1
+XAV_DEBUG_LOGS = -DXAV_DEBUG=1
 
-HEADERS = $(XAV_DIR)/reader.h $(XAV_DIR)/decoder.h $(XAV_DIR)/utils.h
-SOURCES = $(XAV_DIR)/xav_nif.c $(XAV_DIR)/reader.c $(XAV_DIR)/decoder.c $(XAV_DIR)/utils.c
+HEADERS = $(XAV_DIR)/reader.h $(XAV_DIR)/decoder.h $(XAV_DIR)/converter.h $(XAV_DIR)/utils.h
+SOURCES = $(XAV_DIR)/xav_nif.c $(XAV_DIR)/reader.c $(XAV_DIR)/decoder.c $(XAV_DIR)/converter.c $(XAV_DIR)/utils.c
 
 CFLAGS = $(XAV_DEBUG_LOGS) -fPIC -shared
 IFLAGS = -I$(ERTS_INCLUDE_DIR) -I$(XAV_DIR)

c_src/xav/converter.c

@@ -0,0 +1,66 @@
+#include "converter.h"
+#include <libavutil/channel_layout.h>
+#include <libavutil/opt.h>
+#include <libavutil/samplefmt.h>
+#include <libswresample/swresample.h>
+#include <stdint.h>
+
+#include "utils.h"
+
+int converter_init(struct Converter *c, AVChannelLayout in_chlayout, int in_sample_rate,
+                   enum AVSampleFormat in_sample_fmt, AVChannelLayout out_chlayout,
+                   int out_sample_rate, enum AVSampleFormat out_sample_fmt) {
+  c->swr_ctx = swr_alloc();
+  c->in_sample_rate = in_sample_rate;
+  c->out_sample_rate = out_sample_rate;
+  c->out_chlayout = out_chlayout;
+  c->out_sample_fmt = out_sample_fmt;
+
+  av_opt_set_chlayout(c->swr_ctx, "in_chlayout", &in_chlayout, 0);
+  av_opt_set_chlayout(c->swr_ctx, "out_chlayout", &out_chlayout, 0);
+
+  av_opt_set_int(c->swr_ctx, "in_sample_rate", in_sample_rate, 0);
+  av_opt_set_int(c->swr_ctx, "out_sample_rate", out_sample_rate, 0);
+
+  av_opt_set_sample_fmt(c->swr_ctx, "in_sample_fmt", in_sample_fmt, 0);
+  av_opt_set_sample_fmt(c->swr_ctx, "out_sample_fmt", out_sample_fmt, 0);
+
+  return swr_init(c->swr_ctx);
+}
+
+int converter_convert(struct Converter *c, AVFrame *src_frame, uint8_t ***out_data,
+                      int *out_samples, int *out_size) {
+  uint8_t **out_data_tmp = NULL;
+  int max_out_nb_samples = swr_get_out_samples(c->swr_ctx, src_frame->nb_samples);
+  int out_bytes_per_sample = av_get_bytes_per_sample(c->out_sample_fmt);
+
+  // Some parts of ffmpeg require buffers to by divisible by 32
+  // to use fast/aligned SIMD routines - this is what align option is used for.
+  // See https://stackoverflow.com/questions/35678041/what-is-linesize-alignment-meaning
+  // Because we return the binary straight to the Erlang, we can disable it.
+  int ret = av_samples_alloc_array_and_samples(&out_data_tmp, NULL, c->out_chlayout.nb_channels,
+                                               max_out_nb_samples, c->out_sample_fmt, 1);
+
+  if (ret < 0) {
+    XAV_LOG_DEBUG("Couldn't allocate array for out samples.");
+    return ret;
+  }
+
+  *out_data = out_data_tmp;
+
+  *out_samples = swr_convert(c->swr_ctx, out_data_tmp, max_out_nb_samples,
+                             (const uint8_t **)src_frame->data, src_frame->nb_samples);
+
+  if (*out_samples < 0) {
+    XAV_LOG_DEBUG("Couldn't convert samples: %d", *out_samples);
+    return -1;
+  }
+
+  XAV_LOG_DEBUG("Converted %d samples per channel", *out_samples);
+
+  *out_size = *out_samples * out_bytes_per_sample * c->out_chlayout.nb_channels;
+
+  return 0;
+}
+
+void converter_free(struct Converter *c) { swr_free(&c->swr_ctx); }

c_src/xav/converter.h

@@ -0,0 +1,18 @@
+#include <libavutil/channel_layout.h>
+#include <libswresample/swresample.h>
+#include <stdint.h>
+
+struct Converter {
+  SwrContext *swr_ctx;
+  int64_t in_sample_rate;
+  int64_t out_sample_rate;
+  AVChannelLayout out_chlayout;
+  enum AVSampleFormat out_sample_fmt;
+};
+
+int converter_init(struct Converter *c, AVChannelLayout in_chlayout, int in_sample_rate,
+                   enum AVSampleFormat in_sample_fmt, AVChannelLayout out_chlaout,
+                   int out_sample_rate, enum AVSampleFormat out_sample_fmt);
+int converter_convert(struct Converter *c, AVFrame *src_frame, uint8_t ***out_data,
+                      int *out_samples, int *out_size);
+void converter_free(struct Converter *converter);

c_src/xav/decoder.c

@@ -67,10 +67,10 @@ int decoder_decode(struct Decoder *decoder, AVPacket *pkt, AVFrame *frame) {
       }
     }
 
-    if (convert_to_interleaved(decoder->swr_ctx, frame, decoder->rgb_dst_data,
-                               decoder->rgb_dst_linesize) != 0) {
-      return -1;
-    }
+    // if (convert_to_interleaved(decoder->swr_ctx, frame, decoder->rgb_dst_data,
+    //                            decoder->rgb_dst_linesize) != 0) {
+    //   return -1;
+    // }
 
     decoder->frame_data = decoder->rgb_dst_data;
     decoder->frame_linesize = decoder->rgb_dst_linesize;

c_src/xav/reader.c

@@ -1,5 +1,6 @@
 #include "reader.h"
 #include "utils.h"
+#include <libavutil/samplefmt.h>
 #include <libavutil/version.h>
 
 int reader_init(struct Reader *reader, unsigned char *path, size_t path_size, int device_flag,
@@ -69,23 +70,13 @@ int reader_init(struct Reader *reader, unsigned char *path, size_t path_size, in
   }
 
   if (reader->media_type == AVMEDIA_TYPE_AUDIO) {
-    reader->swr_ctx = swr_alloc();
-    enum AVSampleFormat out_sample_fmt = av_get_alt_sample_fmt(reader->c->sample_fmt, 0);
-
-#if LIBAVUTIL_VERSION_MAJOR >= 58
-    av_opt_set_chlayout(reader->swr_ctx, "in_chlayout", &reader->c->ch_layout, 0);
-    av_opt_set_chlayout(reader->swr_ctx, "out_chlayout", &reader->c->ch_layout, 0);
-#else
-    av_opt_set_channel_layout(reader->swr_ctx, "in_channel_layout", reader->c->channel_layout, 0);
-    av_opt_set_channel_layout(reader->swr_ctx, "out_channel_layout", reader->c->channel_layout, 0);
-#endif
-
-    av_opt_set_int(reader->swr_ctx, "in_sample_rate", reader->c->sample_rate, 0);
-    av_opt_set_int(reader->swr_ctx, "out_sample_rate", reader->c->sample_rate, 0);
-    av_opt_set_sample_fmt(reader->swr_ctx, "in_sample_fmt", reader->c->sample_fmt, 0);
-    av_opt_set_sample_fmt(reader->swr_ctx, "out_sample_fmt", out_sample_fmt, 0);
-
-    ret = swr_init(reader->swr_ctx);
+    AVChannelLayout out_chlayout = AV_CHANNEL_LAYOUT_MONO;
+    int out_sample_rate = 16000;
+    enum AVSampleFormat out_sample_fmt = AV_SAMPLE_FMT_FLT;
+
+    int ret = converter_init(&reader->converter, reader->c->ch_layout, reader->c->sample_rate,
+                             reader->c->sample_fmt, out_chlayout, out_sample_rate, out_sample_fmt);
+
     if (ret < 0) {
       return ret;
     }
@@ -156,6 +147,13 @@ int reader_next_frame(struct Reader *reader) {
 
     if (ret == 0) {
       XAV_LOG_DEBUG("Received frame");
+      XAV_LOG_DEBUG("frame->linesize[0] %d", reader->frame->linesize[0]);
+      XAV_LOG_DEBUG("frame->ch_layout.nb_channels %d", reader->frame->ch_layout.nb_channels);
+      XAV_LOG_DEBUG("frame->nb_samples %d", reader->frame->nb_samples);
+      XAV_LOG_DEBUG("av_get_bytes_per_sample(frame->format) %d",
+                    av_get_bytes_per_sample(reader->frame->format));
+      XAV_LOG_DEBUG("av_get_sample_fmt_name(frame->format) %s",
+                    av_get_sample_fmt_name(reader->frame->format));
       frame_ready = 1;
     } else if (ret == AVERROR_EOF) {
       XAV_LOG_DEBUG("EOF");
@@ -200,19 +198,10 @@ int reader_next_frame(struct Reader *reader) {
   } else if (reader->media_type == AVMEDIA_TYPE_VIDEO) {
     reader->frame_data = reader->frame->data;
     reader->frame_linesize = reader->frame->linesize;
-  } else if (reader->media_type == AVMEDIA_TYPE_AUDIO &&
-             av_sample_fmt_is_planar(reader->frame->format) == 1) {
-    XAV_LOG_DEBUG("Converting to interleaved");
-
-    if (convert_to_interleaved(reader->swr_ctx, reader->frame, reader->rgb_dst_data,
-                               reader->rgb_dst_linesize) != 0) {
-      return -1;
-    }
-
-    reader->frame_data = reader->rgb_dst_data;
-    reader->frame_linesize = reader->rgb_dst_linesize;
-  } else {
-    reader->frame_data = reader->frame->extended_data;
+  } else if (reader->media_type == AVMEDIA_TYPE_AUDIO) {
+    XAV_LOG_DEBUG("Converting to out format");
+    return converter_convert(&reader->converter, reader->frame, &reader->out_data,
+                             &reader->out_samples, &reader->out_size);
   }
 
   return 0;
@@ -226,6 +215,7 @@ void reader_free_frame(struct Reader *reader) {
              reader->frame_data == reader->rgb_dst_data) {
     av_freep(&reader->frame_data[0]);
   }
+  free(reader->out_data);
 }
 
 void reader_free(struct Reader *reader) {

c_src/xav/reader.h

@@ -8,8 +8,8 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <time.h>
 
+#include "converter.h"
 #include "utils.h"
 
 struct Reader {
@@ -39,6 +39,16 @@ struct Reader {
   // whether convertion to rgb was needed
   uint8_t **frame_data;
   int *frame_linesize;
+
+  struct Converter converter;
+  // Buffer where audio samples are written after conversion.
+  // We always convet to packed format, so only out_data[0] is set.
+  uint8_t **out_data;
+  // Number of samples in out_data buffer
+  int out_samples;
+  // Size of out_data buffer.
+  // This is the same as out_samples * bytes_per_sample(out_format) * out_channels.
+  int out_size;
 };
 
 int reader_init(struct Reader *reader, unsigned char *path, size_t path_size, int device_flag,

c_src/xav/utils.c

@@ -1,4 +1,7 @@
 #include "utils.h"
+#include <libavutil/mathematics.h>
+#include <libavutil/opt.h>
+#include <stdint.h>
 
 void print_supported_pix_fmts(AVCodec *codec) {
   if (codec->pix_fmts == NULL) {
@@ -42,31 +45,6 @@ void convert_to_rgb(AVFrame *src_frame, uint8_t *dst_data[], int dst_linesize[])
             src_frame->height, dst_data, dst_linesize);
 }
 
-int convert_to_interleaved(SwrContext *swr_ctx, AVFrame *src_frame, uint8_t **dst_data,
-                           int *dst_linesize) {
-#if LIBAVUTIL_VERSION_MAJOR >= 58
-  int channels = src_frame->ch_layout.nb_channels;
-#else
-  int channels = src_frame->channels;
-#endif
-
-  int samples_per_channel = src_frame->nb_samples;
-
-  int ret =
-      av_samples_alloc(dst_data, dst_linesize, channels, samples_per_channel, src_frame->format, 0);
-  if (ret < 0) {
-    return ret;
-  }
-
-  ret = swr_convert(swr_ctx, dst_data, samples_per_channel, (const uint8_t **)src_frame->data,
-                    samples_per_channel);
-  if (ret < 0) {
-    return ret;
-  }
-
-  return 0;
-}
-
 ERL_NIF_TERM xav_nif_ok(ErlNifEnv *env, ERL_NIF_TERM data_term) {
   ERL_NIF_TERM ok_term = enif_make_atom(env, "ok");
   return enif_make_tuple(env, 2, ok_term, data_term);

c_src/xav/utils.h

@@ -20,8 +20,6 @@
 void print_supported_pix_fmts(AVCodec *codec);
 int init_swr_ctx_from_frame(SwrContext **swr_ctx, AVFrame *frame);
 void convert_to_rgb(AVFrame *src_frame, uint8_t *dst_data[], int dst_linesize[]);
-int convert_to_interleaved(SwrContext *swr_ctx, AVFrame *src_frame, uint8_t **dst_data,
-                           int *dst_linesize);
 
 ERL_NIF_TERM xav_nif_ok(ErlNifEnv *env, ERL_NIF_TERM data_term);
 ERL_NIF_TERM xav_nif_error(ErlNifEnv *env, char *reason);

c_src/xav/xav_nif.c

@@ -93,8 +93,21 @@ ERL_NIF_TERM next_frame(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     frame_term = xav_nif_video_frame_to_term(env, reader->frame, reader->frame_data,
                                              reader->frame_linesize, reader->out_format_name);
   } else if (reader->media_type == AVMEDIA_TYPE_AUDIO) {
-    frame_term = xav_nif_audio_frame_to_term(env, reader->frame, reader->frame_data,
-                                             reader->out_format_name);
+    // frame_term = xav_nif_audio_frame_to_term(env, reader->frame, reader->frame_data,
+    //                                          reader->out_format_name);
+
+    ERL_NIF_TERM data_term;
+
+    unsigned char *ptr = enif_make_new_binary(env, reader->out_size, &data_term);
+    memcpy(ptr, reader->out_data[0], reader->out_size);
+
+    const char *out_format_name = av_get_sample_fmt_name(reader->converter.out_sample_fmt);
+
+    ERL_NIF_TERM samples_term = enif_make_int(env, reader->out_samples);
+    ERL_NIF_TERM format_term = enif_make_atom(env, out_format_name);
+    ERL_NIF_TERM pts_term = enif_make_int(env, reader->frame->pts);
+
+    frame_term = enif_make_tuple(env, 4, data_term, format_term, samples_term, pts_term);
   }
 
   reader_free_frame(reader);

lib/frame.ex

@@ -57,7 +57,7 @@ defmodule Xav.Frame do
   end
 
   def to_nx(%__MODULE__{type: :audio} = frame) do
-    Nx.from_binary(frame.data, to_nx_format(frame.format))
+    Nx.from_binary(frame.data, to_nx_format(frame.format), backend: Nx.BinaryBackend)
   end
 
   defp to_nx_format(:u8), do: :u8

test/reader_test.exs

@@ -45,8 +45,53 @@ defmodule Xav.ReaderTest do
   test "speech to text" do
     # This file has been downloaded from https://audio-samples.github.io/
     # Section: Samples from the model without biasing or priming.
-    reader = Xav.Reader.new!("./test/fixtures/melnet_sample_0.mp3", read: :audio)
+    # reader = Xav.Reader.new!("./test/fixtures/melnet_sample_0.mp3", read: :audio)
+    # reader = Xav.Reader.new!("./test/fixtures/harvard.wav", read: :audio)
+    reader = Xav.Reader.new!("./test/fixtures/harvard.wav", read: :audio)
+    # reader = Xav.Reader.new!("./test/fixtures/harvard_converted.mp3", read: :audio)
+    # reader = Xav.Reader.new!("./test/fixtures/harvard_converted.wav", read: :audio)
+    # reader = Xav.Reader.new!("./test/fixtures/harvard_converted.mp3", read: :audio)
 
+    # {:ok, whisper} = Bumblebee.load_model({:hf, "openai/whisper-tiny"})
+    # {:ok, featurizer} = Bumblebee.load_featurizer({:hf, "openai/whisper-tiny"})
+    # {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/whisper-tiny"})
+    # {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "openai/whisper-tiny"})
+
+    # serving =
+    #   Bumblebee.Audio.speech_to_text_whisper(whisper, featurizer, tokenizer, generation_config,
+    #     defn_options: [compiler: EXLA]
+    #   )
+
+    Xav.Reader.next_frame(reader) |> IO.inspect(limit: :infinity)
+
+    # frames = read_frames(reader)
+
+    # IO.inspect(length(frames))
+
+    # batch =
+    #   frames
+    #   |> Enum.filter(fn frame -> frame.data != <<>> end)
+    #   |> Enum.map(&Xav.Frame.to_nx(&1))
+    #   |> Nx.Batch.concatenate()
+
+    # batch = Nx.Defn.jit_apply(&Function.identity/1, [batch])
+    # assert %{chunks: chunks} = Nx.Serving.run(serving, batch)
+
+    # dbg(chunks)
+    # assert chunks != []
+    # assert [
+    #          %{
+    #            text: """
+    #             My thought I have nobody by a beauty and will as you poured. \
+    #            Mr. Rochester has served in that so-done fine-simpless and \
+    #            devoted to bowed, to let might in a\
+    #            """
+    #          }
+    #        ] = chunks
+  end
+
+  @tag :debug2
+  test "" do
     {:ok, whisper} = Bumblebee.load_model({:hf, "openai/whisper-tiny"})
     {:ok, featurizer} = Bumblebee.load_featurizer({:hf, "openai/whisper-tiny"})
     {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/whisper-tiny"})
@@ -57,23 +102,9 @@ defmodule Xav.ReaderTest do
         defn_options: [compiler: EXLA]
       )
 
-    batch =
-      read_frames(reader)
-      |> Enum.map(&Xav.Frame.to_nx(&1))
-      |> Nx.Batch.concatenate()
-
-    batch = Nx.Defn.jit_apply(&Function.identity/1, [batch])
-    assert %{chunks: chunks} = Nx.Serving.run(serving, batch)
-
-    assert [
-             %{
-               text: """
-                My thought I have nobody by a beauty and will as you poured. \
-               Mr. Rochester has served in that so-done fine-simpless and \
-               devoted to bowed, to let might in a\
-               """
-             }
-           ] = chunks
+    output =
+      Nx.Serving.run(serving, {:file, "./test/fixtures/harvard.wav"})
+      |> dbg()
   end
 
   defp read_frames(reader, acc \\ []) do