skyler
/
Lightningbeam
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

gpu-video-encoder: mux the zero-copy encode into a container file 

ZeroCopyEncoder::new now takes an output path and writes a real container
(format inferred from the extension, e.g. .mp4): create an output format
context, add the h264 stream from the encoder, write header; encode_rgba
rescales each packet's ts and av_interleaved_write_frame's it; finish
flushes + writes the trailer + closes. Sets AV_CODEC_FLAG_GLOBAL_HEADER for
mp4/mov so SPS/PPS land in extradata. This lets the editor's existing
mux_video_and_audio consume the temp video file unchanged.

The zerocopy_encode test now writes a .mp4 and ffprobe-verifies the codec,
dimensions, and frame count. Also let wgpu own the imported plane-image
destruction via texture_from_raw drop callbacks (clears two warnings).
This commit is contained in:
Skyler Lehmkuhl 2026-06-25 15:42:59 -04:00
parent ba897eaea2
commit a00e73c4b3
3 changed files with 111 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs
View File

@ -149,13 +149,13 @@ pub fn import_raw(drm: &DrmDevice, buf: &Nv12DmaBuf) -> Result<ImportedNv12, Str
.device
.as_hal::<wgpu_hal::vulkan::Api>()
.ok_or("device is not Vulkan")?;
let mut wrap = |img: vk::Image, format: wgpu::TextureFormat, w: u32, h: u32| -> wgpu::Texture {
let wrap = |img: vk::Image, format: wgpu::TextureFormat, w: u32, h: u32| -> wgpu::Texture {
// wgpu destroys the image (after wait-idle) when the texture drops; the
// captured Arc<MemoryGuard> frees the shared memory once both have run.
let dev = device.clone();
let guard = mem_guard.clone();
let cb: wgpu_hal::DropCallback = Box::new(move || {
unsafe { dev.destroy_image(img, None) };
dev.destroy_image(img, None);
drop(guard);
});
let hal_desc = wgpu_hal::TextureDescriptor {

99
lightningbeam-ui/gpu-video-encoder/src/encoder.rs
View File

@ -10,6 +10,7 @@ use crate::vk_device::{self, DrmDevice};
use ffmpeg_sys_next as ff;
use std::collections::HashMap;
use std::ffi::CString;
use std::path::Path;
use std::ptr;
#[inline]
@ -24,16 +25,26 @@ pub struct ZeroCopyEncoder {
frames_ref: *mut ff::AVBufferRef,
enc: *mut ff::AVCodecContext,
pkt: *mut ff::AVPacket,
/// Output container (e.g. `.mp4`); packets are muxed into it directly.
oc: *mut ff::AVFormatContext,
enc_tb: ff::AVRational,
stream_tb: ff::AVRational,
width: u32,
height: u32,
pts: i64,
cache: HashMap<usize, ImportedNv12>,
out: Vec<u8>,
}
impl ZeroCopyEncoder {
/// Build a zero-copy `h264_vaapi` encoder, or `Err` if VAAPI/the device is unavailable.
pub fn new(width: u32, height: u32, framerate: i32, bitrate_kbps: u32) -> Result<Self, String> {
/// Build a zero-copy `h264_vaapi` encoder writing to `output_path` (container inferred
/// from the extension, e.g. `.mp4`). `Err` if VAAPI/the device is unavailable.
pub fn new(
width: u32,
height: u32,
framerate: i32,
bitrate_kbps: u32,
output_path: &Path,
) -> Result<Self, String> {
let drm = vk_device::create()?;
let renderer = Rgba2Nv12::new(&drm.device);
unsafe {
@ -81,14 +92,60 @@ impl ZeroCopyEncoder {
}
(*enc).hw_frames_ctx = ff::av_buffer_ref(frames_ref);
if ff::avcodec_open2(enc, codec, ptr::null_mut()) < 0 {
// Output container (format inferred from the path's extension).
let cleanup = |frames_ref: *mut ff::AVBufferRef, enc: *mut ff::AVCodecContext, hw: *mut ff::AVBufferRef| {
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::avcodec_free_context(&mut (enc as *mut _));
ff::av_buffer_unref(&mut hw_device);
let mut h = hw;
ff::av_buffer_unref(&mut h);
};
let path_c = CString::new(output_path.to_string_lossy().as_ref()).unwrap();
let mut oc: *mut ff::AVFormatContext = ptr::null_mut();
if ff::avformat_alloc_output_context2(&mut oc, ptr::null(), ptr::null(), path_c.as_ptr()) < 0
|| oc.is_null()
{
cleanup(frames_ref, enc, hw_device);
return Err(format!("avformat_alloc_output_context2 for {output_path:?} failed"));
}
// mp4/mov want SPS/PPS in extradata, not inline — set before opening the encoder.
if (*(*oc).oformat).flags & ff::AVFMT_GLOBALHEADER as i32 != 0 {
(*enc).flags |= ff::AV_CODEC_FLAG_GLOBAL_HEADER as i32;
}
if ff::avcodec_open2(enc, codec, ptr::null_mut()) < 0 {
ff::avformat_free_context(oc);
cleanup(frames_ref, enc, hw_device);
return Err("avcodec_open2(h264_vaapi) failed".into());
}
let stream = ff::avformat_new_stream(oc, codec);
if stream.is_null() {
ff::avformat_free_context(oc);
cleanup(frames_ref, enc, hw_device);
return Err("avformat_new_stream failed".into());
}
if ff::avcodec_parameters_from_context((*stream).codecpar, enc) < 0 {
ff::avformat_free_context(oc);
cleanup(frames_ref, enc, hw_device);
return Err("avcodec_parameters_from_context failed".into());
}
(*stream).time_base = (*enc).time_base;
if ff::avio_open(&mut (*oc).pb, path_c.as_ptr(), ff::AVIO_FLAG_WRITE as i32) < 0 {
ff::avformat_free_context(oc);
cleanup(frames_ref, enc, hw_device);
return Err(format!("avio_open {output_path:?} failed"));
}
if ff::avformat_write_header(oc, ptr::null_mut()) < 0 {
ff::avio_closep(&mut (*oc).pb);
ff::avformat_free_context(oc);
cleanup(frames_ref, enc, hw_device);
return Err("avformat_write_header failed".into());
}
// The muxer may rewrite the stream time_base in write_header.
let stream_tb = (*stream).time_base;
Ok(Self {
drm,
renderer,
@ -96,11 +153,13 @@ impl ZeroCopyEncoder {
frames_ref,
enc,
pkt: ff::av_packet_alloc(),
oc,
enc_tb: (*enc).time_base,
stream_tb,
width,
height,
pts: 0,
cache: HashMap::new(),
out: Vec::new(),
})
}
}
@ -193,20 +252,28 @@ impl ZeroCopyEncoder {
if r < 0 {
return Err(format!("avcodec_receive_packet failed: {r}"));
}
let data = std::slice::from_raw_parts((*self.pkt).data, (*self.pkt).size as usize);
self.out.extend_from_slice(data);
ff::av_packet_unref(self.pkt);
ff::av_packet_rescale_ts(self.pkt, self.enc_tb, self.stream_tb);
(*self.pkt).stream_index = 0;
// Takes ownership of the packet's buffer (unrefs it for us).
let w = ff::av_interleaved_write_frame(self.oc, self.pkt);
if w < 0 {
return Err(format!("av_interleaved_write_frame failed: {w}"));
}
}
Ok(())
}
/// Flush the encoder and return the accumulated Annex-B H.264 bitstream.
pub fn finish(mut self) -> Result<Vec<u8>, String> {
/// Flush the encoder, write the container trailer, and close the output file.
pub fn finish(mut self) -> Result<(), String> {
unsafe {
ff::avcodec_send_frame(self.enc, ptr::null_mut());
self.drain()?;
if ff::av_write_trailer(self.oc) < 0 {
return Err("av_write_trailer failed".into());
}
ff::avio_closep(&mut (*self.oc).pb);
}
Ok(std::mem::take(&mut self.out))
Ok(())
}
}
@ -219,6 +286,14 @@ impl Drop for ZeroCopyEncoder {
let mut fr = self.frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut self.hw_device);
if !self.oc.is_null() {
// `finish` nulls pb via avio_closep; close here too if it wasn't called.
if !(*self.oc).pb.is_null() {
ff::avio_closep(&mut (*self.oc).pb);
}
ff::avformat_free_context(self.oc);
self.oc = ptr::null_mut();
}
}
}
}

43
lightningbeam-ui/gpu-video-encoder/tests/zerocopy_encode.rs
View File

@ -8,7 +8,9 @@ use gpu_video_encoder::encoder::ZeroCopyEncoder;
#[test]
fn zerocopy_encode_h264() {
let (w, h) = (640u32, 480u32);
let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000) {
let out = std::env::temp_dir().join("gpu_video_encoder_zerocopy.mp4");
let _ = std::fs::remove_file(&out);
let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000, &out) {
Ok(e) => e,
Err(e) => {
eprintln!("[zc-encode] unavailable, skipping: {e}");
@ -49,27 +51,26 @@ fn zerocopy_encode_h264() {
enc.encode_rgba(&src).expect("encode_rgba");
}
let h264 = enc.finish().expect("finish");
eprintln!("[zc-encode] {} frames -> {} bytes H.264", n, h264.len());
assert!(h264.len() > 1000, "implausibly small output");
assert!(
h264.starts_with(&[0, 0, 0, 1]) || h264.starts_with(&[0, 0, 1]),
"not Annex-B H.264"
);
enc.finish().expect("finish");
let meta = std::fs::metadata(&out).expect("output .mp4 missing");
eprintln!("[zc-encode] {} frames -> {} bytes mp4 at {}", n, meta.len(), out.display());
assert!(meta.len() > 1000, "implausibly small output");
// Write it out and ffprobe-verify if ffprobe is present.
let out = std::env::temp_dir().join("gpu_video_encoder_zerocopy.h264");
std::fs::write(&out, &h264).unwrap();
eprintln!("[zc-encode] wrote {}", out.display());
if let Ok(o) = std::process::Command::new("ffprobe")
.args(["-hide_banner", "-v", "error", "-show_entries", "stream=codec_name,width,height", "-of", "default=noprint_wrappers=1"])
// ffprobe-verify the container: H.264 stream, right dims, ~n frames.
let o = std::process::Command::new("ffprobe")
.args([
"-hide_banner", "-v", "error", "-count_frames",
"-show_entries", "stream=codec_name,width,height,nb_read_frames",
"-of", "default=noprint_wrappers=1",
])
.arg(&out)
.output()
{
let s = String::from_utf8_lossy(&o.stdout);
eprintln!("[zc-encode] ffprobe:\n{s}");
assert!(s.contains("codec_name=h264"), "ffprobe didn't see H.264");
assert!(s.contains(&format!("width={w}")), "wrong width");
}
eprintln!("[zc-encode] ✅ zero-copy H.264 encode verified");
.expect("run ffprobe");
let s = String::from_utf8_lossy(&o.stdout);
eprintln!("[zc-encode] ffprobe:\n{s}");
assert!(s.contains("codec_name=h264"), "ffprobe didn't see H.264");
assert!(s.contains(&format!("width={w}")), "wrong width");
assert!(s.contains(&format!("height={h}")), "wrong height");
assert!(s.contains(&format!("nb_read_frames={n}")), "expected {n} frames");
eprintln!("[zc-encode] ✅ zero-copy H.264 mp4 encode verified");
}