Add encoder

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
 resolver = "2"
-members = ["native/vkvideo_decoder"]
+members = ["native/vkvideo"]

lib/decoder.ex

@@ -5,10 +5,10 @@ defmodule Membrane.VKVideo.Decoder do
   """
   use Membrane.Filter
 
-  alias __MODULE__.Native
+  alias Membrane.VKVideo.{DeviceServer, Native}
 
   def_input_pad :input, accepted_format: %Membrane.H264{stream_structure: :annexb, alignment: :au}
-  def_output_pad :output, accepted_format: Membrane.RawVideo
+  def_output_pad :output, accepted_format: %Membrane.RawVideo{pixel_format: :NV12}
 
   @impl true
   def handle_init(_ctx, _opts) do
@@ -18,7 +18,8 @@ defmodule Membrane.VKVideo.Decoder do
 
   @impl true
   def handle_setup(_ctx, state) do
-    {:ok, decoder} = Native.new()
+    {:ok, device} = DeviceServer.get_device()
+    {:ok, decoder} = Native.new_decoder(device)
     state = %{state | decoder: decoder}
     {[], state}
   end
@@ -70,8 +71,9 @@ defmodule Membrane.VKVideo.Decoder do
 
   @impl true
   def handle_end_of_stream(:input, _ctx, state) do
-    {:ok, flushed_frames} = Native.flush(state.decoder)
+    {:ok, flushed_frames} = Native.flush_decoder(state.decoder)
+    :ok = Native.destroy(state.decoder)
     {actions, state} = Enum.flat_map_reduce(flushed_frames, state, &prepare_actions(&1, &2))
-    {actions ++ [end_of_stream: :output], state}
+    {actions ++ [end_of_stream: :output], %{state | decoder: nil}}
   end
 end

lib/device_server.ex

@@ -0,0 +1,28 @@
+defmodule Membrane.VKVideo.DeviceServer do
+  @moduledoc false
+  use GenServer
+  alias Membrane.VKVideo.Native
+
+  @impl true
+  def init(_opts) do
+    {:ok, %{device: nil}}
+  end
+
+  @impl true
+  def handle_call(:get_device, _from, state) do
+    state = maybe_create_device(state)
+    {:reply, state.device, state}
+  end
+
+  defp maybe_create_device(%{device: nil} = state) do
+    device = Native.create_device()
+    %{state | device: device}
+  end
+
+  defp maybe_create_device(state), do: state
+
+  @spec get_device() :: {:ok, Native.t()} | no_return()
+  def get_device() do
+    GenServer.call(__MODULE__, :get_device)
+  end
+end

lib/encoder.ex

@@ -0,0 +1,166 @@
+defmodule Membrane.VKVideo.Encoder do
+  @moduledoc """
+  H.264 encoder taking advantage of hardware acceleration provided
+  by Vulkan video extensions.
+  """
+  use Membrane.Filter
+
+  require Membrane.Logger
+
+  alias Membrane.VKVideo.{DeviceServer, Native}
+
+  def_input_pad :input, accepted_format: %Membrane.RawVideo{pixel_format: :NV12}
+
+  def_output_pad :output,
+    accepted_format: %Membrane.H264{stream_structure: :annexb, alignment: :au}
+
+  def_options tune: [
+                spec: :low_latency | :high_quality,
+                default: :low_latency,
+                description: """
+                Specifies whether the encoder should be optimized for minimal latency (which is
+                important in case of livestreams) or for higher quality (applicable to offline encoding).
+                """
+              ],
+              approx_framerate: [
+                spec: {non_neg_integer(), pos_integer()} | nil,
+                default: nil,
+                description: """
+                Framerate of the stream expressed in number of frames per second.
+                It's only used by the rate control mechanism and therefore it does not need to be an exact
+                value. If nil, the framerate will be read from the stream format's structure or set
+                to fixed value of 30 frames per second if framerate is not provided by the stream format.
+                """
+              ],
+              rate_control: [
+                spec:
+                  :encoder_default
+                  | :disabled
+                  | {:variable_bitrate, __MODULE__.VariableBitrate.t()}
+                  | {:constant_bitrate, __MODULE__.ConstantBitrate.t()},
+                default: :encoder_default,
+                description: """
+                Specifies which rate control mechanism should by used by the encoder.
+                """
+              ]
+
+  @impl true
+  def handle_init(_ctx, opts) do
+    state = %{
+      encoder: nil,
+      width: nil,
+      height: nil,
+      override_framerate?: opts.approx_framerate != nil,
+      framerate: opts.approx_framerate,
+      rate_control: opts.rate_control,
+      tune: opts.tune
+    }
+
+    {[], state}
+  end
+
+  @impl true
+  def handle_stream_format(:input, stream_format, _ctx, state) do
+    cond do
+      state.override_framerate? and
+          (stream_format.width != state.width or
+             stream_format.height != state.height) ->
+        %{
+          state
+          | width: stream_format.width,
+            height: stream_format.height
+        }
+        |> spawn_encoder()
+
+      not state.override_framerate? and
+          (stream_format.width != state.width or stream_format.height != state.height or
+             stream_format.framerate != state.framerate) ->
+        %{
+          state
+          | width: stream_format.width,
+            height: stream_format.height,
+            framerate: resolve_framerate(stream_format, state)
+        }
+        |> spawn_encoder()
+
+      true ->
+        {[], state}
+    end
+  end
+
+  defp spawn_encoder(state) do
+    {:ok, device} = DeviceServer.get_device()
+
+    {:ok, encoder} =
+      Native.new_encoder(
+        device,
+        state.width,
+        state.height,
+        state.framerate,
+        state.tune,
+        state.rate_control
+      )
+
+    state = put_in(state, [:encoder], encoder)
+
+    stream_format =
+      %Membrane.H264{
+        stream_structure: :annexb,
+        alignment: :au,
+        width: state.width,
+        height: state.height
+      }
+
+    stream_format =
+      if state.override_framerate? do
+        stream_format
+      else
+        %{stream_format | framerate: state.framerate}
+      end
+
+    {[stream_format: {:output, stream_format}], state}
+  end
+
+  defp resolve_framerate(stream_format, state) do
+    if is_nil(stream_format.framerate) and requires_framerate?(state.rate_control) do
+      Membrane.Logger.warning("""
+          Framerate is required when using #{inspect(elem(state.rate_control, 0))} rate control but it
+          wasn't provided in the stream format nor via options. Please provide approximate framerate
+          using `approx_framerate` option of the element.
+      """)
+    end
+
+    stream_format.framerate || {30, 1}
+  end
+
+  defp requires_framerate?({:constant_bitrate, _constant_bitrate}), do: true
+  defp requires_framerate?({:variable_bitrate, _variable_bitrate}), do: true
+  defp requires_framerate?(_other_rate_control), do: false
+
+  @impl true
+  def handle_buffer(:input, buffer, _ctx, state) do
+    {:ok, encoded_frame} = Native.encode(state.encoder, buffer.payload, buffer.pts)
+
+    pts =
+      if encoded_frame.pts_ns != nil,
+        do: Membrane.Time.nanoseconds(encoded_frame.pts_ns),
+        else: nil
+
+    {[
+       buffer:
+         {:output,
+          %Membrane.Buffer{
+            payload: encoded_frame.payload,
+            pts: pts,
+            dts: buffer.dts
+          }}
+     ], state}
+  end
+
+  @impl true
+  def handle_end_of_stream(:input, _ctx, state) do
+    :ok = Native.destroy(state.encoder)
+    state = %{state | encoder: nil}
+    {[end_of_stream: :output], state}
+  end
+end

lib/encoder/constant_bitrate.ex

@@ -0,0 +1,12 @@
+defmodule Membrane.VKVideo.Encoder.ConstantBitrate do
+  @moduledoc """
+  Defines encoder setting for constant bitrate rate control algorithm.
+  The following fields need to be specified:
+  * bitrate - desired bitrate of the stream; expressed in bits per second.
+  * virtual_buffer_size_ms - virtual buffer duration for rate control smoothing; larger values increase bitrate stability, smaller values improve responsiveness to scene changes; expressed in milliseconds, defaults to 2 seconds.
+  """
+
+  @type t :: %__MODULE__{bitrate: non_neg_integer(), virtual_buffer_size_ms: non_neg_integer()}
+  @enforce_keys [:bitrate]
+  defstruct @enforce_keys ++ [virtual_buffer_size_ms: 2000]
+end

lib/encoder/variable_bitrate.ex

@@ -0,0 +1,20 @@
+defmodule Membrane.VKVideo.Encoder.VariableBitrate do
+  @moduledoc """
+  Defines encoder setting for variable bitrate rate control algorithm.
+
+  The following fields need to be specified:
+  * average_bitrate - Target average bitrate for VBR encoding; the encoder will try to meet this
+  average over the sequence; expressed in bits per second.
+  * max_bitrate - Maximum allowed bitrate in VBR encoding; caps peak bitrate to prevent excessive
+  spikes while maintaining average bitrate constraints; expressed in bits per second.
+  * virtual_buffer_size_ms - virtual buffer duration for rate control smoothing; larger values increase bitrate stability, smaller values improve responsiveness to scene changes; expressed in milliseconds, defaults to 2 seconds.
+  """
+
+  @type t :: %__MODULE__{
+          average_bitrate: non_neg_integer(),
+          max_bitrate: non_neg_integer(),
+          virtual_buffer_size_ms: non_neg_integer()
+        }
+  @enforce_keys [:average_bitrate, :max_bitrate, :virtual_buffer_size_ms]
+  defstruct @enforce_keys
+end

lib/native.ex

@@ -0,0 +1,52 @@
+defmodule Membrane.VKVideo.Native do
+  @moduledoc false
+  use Rustler, otp_app: :membrane_vk_video_plugin, crate: :vkvideo
+
+  @type t :: reference()
+  @type raw_frame :: %{
+          payload: binary(),
+          pts: non_neg_integer() | nil,
+          width: non_neg_integer(),
+          height: non_neg_integer()
+        }
+
+  @type encoded_frame :: %{
+          payload: binary(),
+          pts: non_neg_integer() | nil
+        }
+
+  @spec create_device() :: {:ok, t()} | no_return()
+  def create_device(), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec new_decoder(t()) :: {:ok, t()} | no_return()
+  def new_decoder(_device), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec decode(t(), binary(), pts_ns :: non_neg_integer() | nil) ::
+          {:ok, raw_frame()} | no_return()
+  def decode(_decoder, _frame, _pts \\ nil), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec flush_decoder(t()) ::
+          {:ok, raw_frame()} | no_return()
+  def flush_decoder(_decoder), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec new_encoder(
+          t(),
+          non_neg_integer(),
+          non_neg_integer(),
+          {non_neg_integer(), non_neg_integer()},
+          :low_latency | :high_quality,
+          :encoder_default
+          | :disabled
+          | {:variable_bitrate, Membrane.VKVideo.Encoder.VariableBitrate.t()}
+          | {:constant_bitrate, Membrane.VKVideo.Encoder.ConstantBitrate.t()}
+        ) :: {:ok, t()} | no_return()
+  def new_encoder(_device, _width, _height, _framerate, _tune, _rate_control),
+    do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec encode(t(), binary(), pts_ns :: non_neg_integer() | nil) ::
+          {:ok, encoded_frame()} | no_return()
+  def encode(_encoder, _raw_frame, _pts \\ nil), do: :erlang.nif_error(:nif_not_loaded)
+
+  @spec destroy(t()) :: :ok
+  def destroy(_resource), do: :erlang.nif_error(:nif_not_loaded)
+end

lib/vk_video.ex

@@ -0,0 +1,11 @@
+defmodule Membrane.VKVideo do
+  @moduledoc false
+  use Application
+
+  alias Membrane.VKVideo.DeviceServer
+
+  @impl true
+  def start(_start_type, _start_args) do
+    GenServer.start_link(DeviceServer, nil, name: DeviceServer)
+  end
+end

mix.exs

@@ -28,7 +28,8 @@ defmodule Membrane.VKVideo.Mixfile do
 
   def application do
     [
-      extra_applications: []
+      extra_applications: [],
+      mod: {Membrane.VKVideo, []}
     ]
   end
 

native/vkvideo_decoder/Cargo.toml → native/vkvideo/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "vkvideo_decoder"
+name = "vkvideo"
 version = "0.1.0"
 authors = []
 edition = "2021"