diff --git a/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs b/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs index 19d17a7..a6bfc33 100644 --- a/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs +++ b/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs @@ -149,13 +149,13 @@ pub fn import_raw(drm: &DrmDevice, buf: &Nv12DmaBuf) -> Result() .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 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 { diff --git a/lightningbeam-ui/gpu-video-encoder/src/encoder.rs b/lightningbeam-ui/gpu-video-encoder/src/encoder.rs index e43ff1b..c209e9b 100644 --- a/lightningbeam-ui/gpu-video-encoder/src/encoder.rs +++ b/lightningbeam-ui/gpu-video-encoder/src/encoder.rs @@ -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, - out: Vec, } 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 { + /// 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 { 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, 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(); + } } } } diff --git a/lightningbeam-ui/gpu-video-encoder/tests/zerocopy_encode.rs b/lightningbeam-ui/gpu-video-encoder/tests/zerocopy_encode.rs index f175d11..69e8ca9 100644 --- a/lightningbeam-ui/gpu-video-encoder/tests/zerocopy_encode.rs +++ b/lightningbeam-ui/gpu-video-encoder/tests/zerocopy_encode.rs @@ -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"); }