diff --git a/Makefile b/Makefile index d45ecb4..d538d2c 100644 --- a/Makefile +++ b/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) diff --git a/c_src/xav/converter.c b/c_src/xav/converter.c new file mode 100644 index 0000000..cecaf99 --- /dev/null +++ b/c_src/xav/converter.c @@ -0,0 +1,66 @@ +#include "converter.h" +#include +#include +#include +#include +#include + +#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); } \ No newline at end of file diff --git a/c_src/xav/converter.h b/c_src/xav/converter.h new file mode 100644 index 0000000..9a574fd --- /dev/null +++ b/c_src/xav/converter.h @@ -0,0 +1,18 @@ +#include +#include +#include + +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); \ No newline at end of file diff --git a/c_src/xav/decoder.c b/c_src/xav/decoder.c index a949ede..a393170 100644 --- a/c_src/xav/decoder.c +++ b/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; diff --git a/c_src/xav/reader.c b/c_src/xav/reader.c index 9d8998b..5fc5412 100644 --- a/c_src/xav/reader.c +++ b/c_src/xav/reader.c @@ -1,5 +1,6 @@ #include "reader.h" #include "utils.h" +#include #include 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) { diff --git a/c_src/xav/reader.h b/c_src/xav/reader.h index 6a7a31d..3522502 100644 --- a/c_src/xav/reader.h +++ b/c_src/xav/reader.h @@ -8,8 +8,8 @@ #include #include #include -#include +#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, diff --git a/c_src/xav/utils.c b/c_src/xav/utils.c index 3cc0c02..9724d79 100644 --- a/c_src/xav/utils.c +++ b/c_src/xav/utils.c @@ -1,4 +1,7 @@ #include "utils.h" +#include +#include +#include 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); diff --git a/c_src/xav/utils.h b/c_src/xav/utils.h index 0e9ca0f..3ddeb22 100644 --- a/c_src/xav/utils.h +++ b/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); diff --git a/c_src/xav/xav_nif.c b/c_src/xav/xav_nif.c index be48433..b58910f 100644 --- a/c_src/xav/xav_nif.c +++ b/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); diff --git a/lib/frame.ex b/lib/frame.ex index d2424a5..f91b3be 100644 --- a/lib/frame.ex +++ b/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 diff --git a/test/reader_test.exs b/test/reader_test.exs index 7ebc5d4..2d86c23 100644 --- a/test/reader_test.exs +++ b/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