Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use xav instead of ffmpeg #403

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

use xav instead of ffmpeg #403

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1b6e128
use xav instead of ffmpeg
kevinschweikert Oct 22, 2024
7b1c02f
use default backend
kevinschweikert Oct 23, 2024
fd81786
transform stream to buffer frame binaries
kevinschweikert Oct 24, 2024
fc45ff5
bring back ffmpeg
kevinschweikert Oct 24, 2024
d882343
make xav optional
kevinschweikert Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 21 additions & 43 deletions lib/bumblebee/audio/speech_to_text_whisper.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ defmodule Bumblebee.Audio.SpeechToTextWhisper do
{:ok, [Nx.backend_transfer(input, Nx.BinaryBackend)]}

{:file, path} when is_binary(path) ->
ffmpeg_read_as_pcm(path, sampling_rate)
from_file(path, sampling_rate)

other ->
cond do
Expand Down Expand Up @@ -164,49 +164,27 @@ defmodule Bumblebee.Audio.SpeechToTextWhisper do
end
end

defp ffmpeg_read_as_pcm(path, sampling_rate) do
channels = 1
defp from_file(path, sampling_rate) do
# This chunk can be of arbitrary size, the serving accumulates
# and overlaps chunks internally as needed.

format =
case System.endianness() do
:little -> "f32le"
:big -> "f32be"
end

cond do
System.find_executable("ffmpeg") == nil ->
{:error, "ffmpeg not found in PATH"}

not File.exists?(path) ->
{:error, "no file found at #{path}"}

true ->
# This chunk can be of arbitrary size, the serving accumulates
# and overlaps chunks internally as needed. We read the file
# as stream to reduce memory usage
chunk_size = 30

stream =
Stream.iterate(0, fn offset -> offset + chunk_size end)
|> Stream.transform({}, fn offset, acc ->
System.cmd(
"ffmpeg",
~w[-ss #{offset} -t #{chunk_size} -i #{path} -ac #{channels} -ar #{sampling_rate} -f #{format} -hide_banner -loglevel quiet pipe:1]
)
|> case do
{<<>>, 0} ->
{:halt, acc}

{data, 0} ->
chunk = Nx.from_binary(data, :f32, backend: Nx.BinaryBackend)
{[chunk], acc}

{_, 1} ->
raise "ffmpeg failed to decode the given file"
end
end)

{:ok, stream}
if File.exists?(path) do
stream =
path
|> Xav.Reader.stream!(
read: :audio,
out_format: :f32,
out_channels: 1,
out_sample_rate: sampling_rate
)
|> Stream.map(fn frame -> Xav.Frame.to_nx(frame) end)
kevinschweikert marked this conversation as resolved.
Show resolved Hide resolved
|> Stream.chunk_every(1000)
|> Stream.map(&Nx.Batch.concatenate/1)
|> Stream.map(fn batch -> Nx.Defn.jit_apply(&Function.identity/1, [batch]) end)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function just needs to return a stream of chunks, so we don't need to do this concatenation.

Suggested change
|> Stream.chunk_every(1000)
|> Stream.map(&Nx.Batch.concatenate/1)
|> Stream.map(fn batch -> Nx.Defn.jit_apply(&Function.identity/1, [batch]) end)

Do you know what determines the length of each chunk, could that be configurable perhaps?

Copy link
Author

@kevinschweikert kevinschweikert Oct 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not completely sure but i'm guessing that a chunk/frame in a video context is the audio duration of one frame. When i remove the Stream.chunk_every/2 pipeline i get a massive performance degradation and the processing does not finish in a reasonable time. I imagine that it's far more efficient to read a chunk of frames first, before sending it to the Whisper serving.
I just reused the implementation of the Elixir WebRTC team from https://github.com/elixir-webrtc/xav/blob/master/README.md

{: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]
  )

# Read a couple of frames.
# See https://hexdocs.pm/bumblebee/Bumblebee.Audio.WhisperFeaturizer.html for default sampling rate.
frames =
    Xav.Reader.stream!("sample.mp3", read: :audio, out_format: :f32, out_channels: 1, out_sample_rate: 16_000)
    |> Stream.take(200)
    |> Enum.map(fn frame -> Xav.Reader.to_nx(frame) end)

batch = Nx.Batch.concatenate(frames)
batch = Nx.Defn.jit_apply(&Function.identity/1, [batch])
Nx.Serving.run(serving, batch)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. The serving transforms the stream to accumulate smaller chunks, but there is a place where we need to append to a list and that may be the reason why it's inefficient with tiny chunks.

However, either way, I think it's wasteful to convert every frame to a tensor just to concatenate later. With the current ffmpeg code we get a single binary for the whole 30s and create a tensor from that. So ideally we want to replicate this.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this suffice?

path
|> Xav.Reader.stream!(
  read: :audio,
  out_format: :f32,
  out_channels: 1,
  out_sample_rate: sampling_rate
)
|> Stream.chunk_every(1000)
|> Stream.map(fn frames ->
  [frame | _] = frames
  binary = Enum.reduce(frames, <<>>, fn frame, acc -> acc <> frame.data end)

  Nx.with_default_backend(Nx.BinaryBackend, fn -> Nx.from_binary(binary, frame.format) end)
end)

The 1000 chunks is currently arbitrary because we don't know the frame_size of the used codec. This could be added to xav i think

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correction: The information is already there. Xav.Frame has a samples field from which we could calculate the 30s chunks which would be calculated by

round(sampling_rate / frame.samples * 30)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handling the frame binaries directly is a good call. I think we can transform the stream, so that we can accumulate the binary, instead of waiting for 1000 binaries and joining then. I was thinking we can determine the number of samples from binary size, but frame.samples is even more convenient.

So it would be something like this:

chunk_samples = sampling_rate * 30

path
|> Xav.Reader.stream!(
  read: :audio,
  out_format: :f32,
  out_channels: 1,
  out_sample_rate: sampling_rate
)
|> Stream.transform(
  fn -> {<<>>, 0} end,
  fn frame, {buffer, samples} ->
    buffer = buffer <> frame.data
    samples = samples + frame.samples

    if samples >= chunk_samples do
      chunk = Nx.from_binary(buffer, :f32, backend: Nx.BinaryBackend)
      {[chunk], {<<>>, 0}}
    else
      {[], {buffer, samples}}
    end
  end,
  fn {buffer, _samples} ->
    chunk = Nx.from_binary(buffer, :f32, backend: Nx.BinaryBackend)
    {[chunk], {<<>>, 0}}
  end,
  fn _ -> :ok end
)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wow! That is a great solution! Thanks!


{:ok, stream}
else
{:error, "no file found at #{path}"}
end
end

Expand Down
5 changes: 3 additions & 2 deletions mix.exs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ defmodule Bumblebee.MixProject do
{:axon, "~> 0.7.0"},
# {:axon, github: "elixir-nx/axon", override: true},
{:tokenizers, "~> 0.4"},
{:nx, "~> 0.9.0"},
{:nx, "~> 0.9.0", override: true},
{:exla, ">= 0.0.0", only: [:dev, :test]},
{:torchx, ">= 0.0.0", only: [:dev, :test]},
# {:nx, github: "elixir-nx/nx", sparse: "nx", override: true},
Expand All @@ -49,7 +49,8 @@ defmodule Bumblebee.MixProject do
{:stb_image, "~> 0.6.0", only: :test},
{:bypass, "~> 2.1", only: :test},
{:ex_doc, "~> 0.28", only: :dev, runtime: false},
{:nx_signal, "~> 0.2.0"}
{:nx_signal, "~> 0.2.0"},
{:xav, "~> 0.6.0"}
]
end

Expand Down
1 change: 1 addition & 0 deletions mix.lock
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,5 +36,6 @@
"torchx": {:hex, :torchx, "0.9.0", "936cbd32233f89d73700c39b7ef56f94b3f3541db03c90f8ddf6b3fe73260e28", [:mix], [{:nx, "~> 0.9.0", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "4e057d6b93fc91191957230b2c61c408861b888abdf6a900baf0db4125405505"},
"unpickler": {:hex, :unpickler, "0.1.0", "c2262c0819e6985b761e7107546cef96a485f401816be5304a65fdd200d5bd6a", [:mix], [], "hexpm", "e2b3f61e62406187ac52afead8a63bfb4e49394028993f3c4c42712743cab79e"},
"unzip": {:hex, :unzip, "0.12.0", "beed92238724732418b41eba77dcb7f51e235b707406c05b1732a3052d1c0f36", [:mix], [], "hexpm", "95655b72db368e5a84951f0bed586ac053b55ee3815fd96062fce10ce4fc998d"},
"xav": {:hex, :xav, "0.6.0", "38835d735fc3d620e41c84fe29cd7db0381436b54c9ef209ba9112255a091fc4", [:make, :mix], [{:elixir_make, "~> 0.7", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:nx, "~> 0.7.0", [hex: :nx, repo: "hexpm", optional: true]}], "hexpm", "ad988df9d44c0ca3ccc4053ea0f1f1702ca14d4b926047b224deb527f0866edf"},
"xla": {:hex, :xla, "0.8.0", "fef314d085dd3ee16a0816c095239938f80769150e15db16dfaa435553d7cb16", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "739c61c8d93b97e12ba0369d10e76130224c208f1a76ad293e3581f056833e57"},
}