From 807ce84be7a2534b7389b622248a0a028f9e4d88 Mon Sep 17 00:00:00 2001 From: modeco80 Date: Sun, 6 Oct 2024 04:33:36 -0400 Subject: [PATCH] server/video: hook up new ctor for nvenc --- server/Cargo.toml | 2 +- server/src/video/encoder_thread.rs | 12 +++++-- server/src/video/h264_encoder.rs | 56 ++++++++++++++++++++++++++++-- server/src/video/hwdevice.rs | 4 +-- server/src/video/hwframe.rs | 4 +-- 5 files changed, 68 insertions(+), 10 deletions(-) diff --git a/server/Cargo.toml b/server/Cargo.toml index 03f3095..46b7417 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -23,7 +23,7 @@ ffmpeg = { version = "7.0.0", package = "ffmpeg-next" } rand = "0.8.5" serde = "1.0.209" serde_json = "1.0.128" -cudarc = "0.12.1" +cudarc = { version = "0.12.1", features = [ "cuda-11050" ] } [patch.crates-io] diff --git a/server/src/video/encoder_thread.rs b/server/src/video/encoder_thread.rs index cffc082..05071ea 100644 --- a/server/src/video/encoder_thread.rs +++ b/server/src/video/encoder_thread.rs @@ -76,6 +76,7 @@ fn encoder_thread_main( let mut sws = None; let mut yuv_frame = None; + //let mut hw_frame = None; let mut frame_number = 0usize; let mut force_keyframe = false; @@ -140,10 +141,15 @@ fn encoder_thread_main( (*mut_yuv_frame.as_mut_ptr()).pts = frame_number as i64; } - enc.send_frame(mut_yuv_frame); + if enc.is_hardware() { - enc.receive_packet(&mut packet) - .expect("failed to recv packet"); + } else { + + enc.send_frame(mut_yuv_frame); + + enc.receive_packet(&mut packet) + .expect("failed to recv packet"); + } // If a packet was recieved dump it unsafe { diff --git a/server/src/video/h264_encoder.rs b/server/src/video/h264_encoder.rs index 4c4e7d9..f192f21 100644 --- a/server/src/video/h264_encoder.rs +++ b/server/src/video/h264_encoder.rs @@ -1,6 +1,7 @@ use super::ffmpeg; use super::hwframe::HwFrameContext; use anyhow::Context; +use cudarc::driver::CudaDevice; use ffmpeg::error::EAGAIN; use ffmpeg::{codec as lavc, packet}; // lavc @@ -22,7 +23,7 @@ pub fn create_context_from_codec(codec: ffmpeg::Codec) -> Result anyhow::Result<(ffmpeg::Codec, ffmpeg::encoder::video::Video)> { @@ -77,7 +78,7 @@ impl H264Encoder { pub fn new_software(size: Size, max_framerate: u32, bitrate: usize) -> anyhow::Result { // Create the libx264 context let (encoder, mut video_encoder_context) = - create_context_and_set_common_parameters("libx264", size, max_framerate, bitrate)?; + create_context_and_set_common_parameters("libx264", &size, max_framerate, bitrate)?; video_encoder_context.set_format(ffmpeg::format::Pixel::YUV420P); @@ -116,6 +117,57 @@ impl H264Encoder { Ok(Self::Software { encoder: encoder }) } + /// Creates a new hardware encoder. + pub fn new_nvenc( + cuda_device: &CudaDevice, + size: Size, + max_framerate: u32, + bitrate: usize + ) -> anyhow::Result { + let (mut encoder, mut video_encoder_context) = + create_context_and_set_common_parameters("h264_nvenc", &size, max_framerate, bitrate) + .with_context(|| "while trying to create encoder")?; + + let cuda_device_context = super::hwdevice::CudaDeviceContextBuilder::new()? + .set_cuda_context((*cuda_device.cu_primary_ctx()) as *mut _) + .build() + .with_context(|| "while trying to create CUDA device context")?; + + 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_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); + + 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; + } + + // set h264_nvenc options + let mut dict = ffmpeg::Dictionary::new(); + + dict.set("profile", "main"); + + // TODO: + dict.set("crf", "43"); + dict.set("crf_max", "48"); + + dict.set("forced-idr", "1"); + + let encoder = video_encoder_context + .open_as_with(encoder, dict) + .with_context(|| "While opening h264_nvenc video codec")?; + + Ok(Self::Nvenc { encoder: encoder, hw_context: hw_frame_context }) + } + // FIXME: It's a bit pointless to have this have a mut borrow, // but you'll probably have a mutable borrow on this already.. pub fn is_hardware(&mut self) -> bool { diff --git a/server/src/video/hwdevice.rs b/server/src/video/hwdevice.rs index ce9a208..ecac8b7 100644 --- a/server/src/video/hwdevice.rs +++ b/server/src/video/hwdevice.rs @@ -35,10 +35,10 @@ pub struct CudaDeviceContextBuilder { } impl CudaDeviceContextBuilder { - pub fn new() -> Result { + pub fn new() -> anyhow::Result { let buffer = unsafe { ffmpeg::sys::av_hwdevice_ctx_alloc(ffmpeg::sys::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA) }; if buffer.is_null() { - return Err("could not allocate a hwdevice".to_string()); + return Err(anyhow::anyhow!("could not allocate a hwdevice".to_string())); } Ok(Self { buffer }) diff --git a/server/src/video/hwframe.rs b/server/src/video/hwframe.rs index e040b6c..b642a87 100644 --- a/server/src/video/hwframe.rs +++ b/server/src/video/hwframe.rs @@ -41,10 +41,10 @@ pub struct HwFrameContextBuilder { } impl HwFrameContextBuilder { - pub fn new(mut cuda_device_context: CudaDeviceContext) -> Result { + pub fn new(mut cuda_device_context: CudaDeviceContext) -> anyhow::Result { let buffer = unsafe { ffmpeg::sys::av_hwframe_ctx_alloc(cuda_device_context.as_raw_mut()) }; if buffer.is_null() { - return Err("could not allocate a hwframe".to_string()); + return Err(anyhow::anyhow!("could not allocate a hwframe")); } Ok(Self { cuda_device_context, buffer })