From fb3f589fd084756dff3c6d200190882c57fa17b0 Mon Sep 17 00:00:00 2001 From: modeco80 Date: Sun, 6 Oct 2024 09:30:17 -0400 Subject: [PATCH] server: working NVENC encoding now I just need to figure out how the hell to reimplement sw --- server/src/main.rs | 3 ++ server/src/video/encoder_thread.rs | 64 +++++++++++++----------------- server/src/video/h264_encoder.rs | 38 +++++++++++++----- 3 files changed, 59 insertions(+), 46 deletions(-) diff --git a/server/src/main.rs b/server/src/main.rs index 2e13ff7..d0201eb 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -143,11 +143,14 @@ async fn main() -> anyhow::Result<()> { // make a new frame for the encoder { let mut lk_frame = frame_clone.lock().expect("Couldn't lock frame"); + + *lk_frame = Some(ffmpeg::frame::Video::new( ffmpeg::format::Pixel::BGRA, size.clone().width, size.clone().height, )); + } let _ = encoder_tx_clone diff --git a/server/src/video/encoder_thread.rs b/server/src/video/encoder_thread.rs index 05071ea..49f11d1 100644 --- a/server/src/video/encoder_thread.rs +++ b/server/src/video/encoder_thread.rs @@ -69,14 +69,12 @@ fn encoder_thread_main( mut rx: mpsc::Receiver, tx: mpsc::Sender, frame: &Arc>>, -) { +) -> anyhow::Result<()> { let mut packet = ffmpeg::Packet::empty(); let mut encoder: Option = None; - let mut sws = None; - let mut yuv_frame = None; - //let mut hw_frame = None; + let dev = cudarc::driver::CudaDevice::new(0)?; let mut frame_number = 0usize; let mut force_keyframe = false; @@ -91,26 +89,22 @@ fn encoder_thread_main( force_keyframe = false; } - yuv_frame = Some(ffmpeg::frame::Video::new( - ffmpeg::format::Pixel::YUV420P, - size.clone().width, - size.clone().height, - )); + - sws = Some( - ffmpeg::software::converter( - size.clone().into(), - ffmpeg::format::Pixel::BGRA, - ffmpeg::format::Pixel::YUV420P, - ) - .expect("Failed to create SWS conversion context"), - ); + encoder = Some(H264Encoder::new_nvenc( + &dev, + size.clone(), + 60, + 5 * (1000 * 1000), + )?); - // TODO: HW support! - encoder = Some( - H264Encoder::new_software(size, 60, 3 * (1000 * 1000)) - .expect("Failed to create encoder"), - ); + + let mut producer_frame_locked = frame.lock().expect("Couldn't lock producer frame"); + let mut producer_frame = producer_frame_locked.as_mut().expect("NOOOO"); + + unsafe { + (*producer_frame.as_mut_ptr()).hw_frames_ctx = encoder.as_mut().unwrap().get_hw_context().as_raw_mut(); + } } EncodeThreadInput::ForceKeyframe => { @@ -122,33 +116,29 @@ fn encoder_thread_main( let enc = encoder.as_mut().unwrap(); // let's encode a frame - let producer_frame_locked = frame.lock().expect("Couldn't lock producer frame"); + let mut producer_frame_locked = frame.lock().expect("Couldn't lock producer frame"); - let producer_frame = producer_frame_locked.as_ref().expect("NOOOO"); - let mut_yuv_frame = yuv_frame.as_mut().unwrap(); + let producer_frame = producer_frame_locked.as_mut().expect("NOOOO"); - let sws_mut = sws.as_mut().unwrap(); - - // scale - sws_mut - .run(producer_frame, mut_yuv_frame) - .expect("Failed to convert producer frame to YUV"); // set the right flags!! - set_frame_flags(mut_yuv_frame, force_keyframe); + set_frame_flags(producer_frame, force_keyframe); unsafe { - (*mut_yuv_frame.as_mut_ptr()).pts = frame_number as i64; + (*producer_frame.as_mut_ptr()).pts = frame_number as i64; + } if enc.is_hardware() { + // should always be Some if we get here on this path + // let mut mut_hw_frame = hw_frame.as_mut().unwrap(); - } else { - - enc.send_frame(mut_yuv_frame); + enc.send_frame(&producer_frame); enc.receive_packet(&mut packet) .expect("failed to recv packet"); + } else { + todo!("FIXME: re-implement SW support."); } // If a packet was recieved dump it @@ -190,6 +180,8 @@ fn encoder_thread_main( } } } + + Ok(()) } pub fn encoder_thread_spawn( diff --git a/server/src/video/h264_encoder.rs b/server/src/video/h264_encoder.rs index f192f21..7dfeb3e 100644 --- a/server/src/video/h264_encoder.rs +++ b/server/src/video/h264_encoder.rs @@ -6,6 +6,7 @@ use ffmpeg::error::EAGAIN; use ffmpeg::{codec as lavc, packet}; // lavc + use crate::types::Size; /// this is required for libx264 to like. Work @@ -40,8 +41,8 @@ fn create_context_and_set_common_parameters( //video_encoder_context.set_max_bit_rate(bitrate); // qp TODO: - //video_encoder_context.set_qmax(30); - //video_encoder_context.set_qmin(35); + video_encoder_context.set_qmax(30); + video_encoder_context.set_qmin(35); video_encoder_context.set_time_base(ffmpeg::Rational(1, max_framerate as i32).invert()); video_encoder_context.set_format(ffmpeg::format::Pixel::YUV420P); @@ -110,6 +111,7 @@ impl H264Encoder { dict.set("forced-idr", "1"); + let encoder = video_encoder_context .open_as_with(encoder, dict) .with_context(|| "While opening x264 video codec")?; @@ -136,23 +138,26 @@ impl H264Encoder { let mut hw_frame_context = super::hwframe::HwFrameContextBuilder::new(cuda_device_context)? .set_width(size.width) .set_height(size.height) - .set_sw_format(ffmpeg::format::Pixel::YUV420P) + .set_sw_format(ffmpeg::format::Pixel::ZRGB32) .set_format(ffmpeg::format::Pixel::CUDA) .build() .with_context(|| "while trying to create CUDA frame context")?; // lol you do not need unsafe code to set this my guy. - video_encoder_context.set_format(ffmpeg::format::Pixel::CUDA); + video_encoder_context.set_format(ffmpeg::format::Pixel::ZRGB32); unsafe { - (*video_encoder_context.as_mut_ptr()).hw_frames_ctx = hw_frame_context.as_raw_mut(); - (*video_encoder_context.as_mut_ptr()).delay = 0; - (*video_encoder_context.as_mut_ptr()).refs = 0; + //(*video_encoder_context.as_mut_ptr()).hw_frames_ctx = hw_frame_context.as_raw_mut(); + //(*video_encoder_context.as_mut_ptr()).delay = 0; + //(*video_encoder_context.as_mut_ptr()).refs = 0; } // set h264_nvenc options let mut dict = ffmpeg::Dictionary::new(); + dict.set("tune", "ull"); + dict.set("preset", "p1"); + dict.set("profile", "main"); // TODO: @@ -160,6 +165,9 @@ impl H264Encoder { dict.set("crf_max", "48"); dict.set("forced-idr", "1"); + + // damn you + dict.set("delay", "0"); let encoder = video_encoder_context .open_as_with(encoder, dict) @@ -177,6 +185,13 @@ impl H264Encoder { } } + pub fn get_hw_context(&mut self) -> &mut HwFrameContext { + match self { + Self::Nvenc { encoder: _, hw_context } => hw_context, + _ => panic!("should not use H264Encoder::get_hw_context() on a Software encoder") + } + } + pub fn send_frame(&mut self, frame: &ffmpeg::Frame) { match self { Self::Software { encoder } => { @@ -188,7 +203,8 @@ impl H264Encoder { hw_context, } => { // Realistically this should be the same right? - todo!("Requires support."); + encoder.send_frame(frame).unwrap(); + //todo!("Requires support."); } } } @@ -204,7 +220,8 @@ impl H264Encoder { hw_context, } => { // Realistically this should be the same right? - todo!("Requires support."); + encoder.send_eof().unwrap(); + // todo!("Requires support."); } } } @@ -218,7 +235,8 @@ impl H264Encoder { hw_context, } => { // Realistically this should be the same right? - todo!("Requires support."); + encoder.receive_packet(packet) + //todo!("Requires support."); } }; }