server: working NVENC encoding

now I just need to figure out how the hell to reimplement sw
This commit is contained in:
Lily Tsuru 2024-10-06 09:30:17 -04:00
parent 807ce84be7
commit fb3f589fd0
3 changed files with 59 additions and 46 deletions

View file

@ -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

View file

@ -69,14 +69,12 @@ fn encoder_thread_main(
mut rx: mpsc::Receiver<EncodeThreadInput>,
tx: mpsc::Sender<EncodeThreadOutput>,
frame: &Arc<Mutex<Option<ffmpeg::frame::Video>>>,
) {
) -> anyhow::Result<()> {
let mut packet = ffmpeg::Packet::empty();
let mut encoder: Option<H264Encoder> = 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"),
);
// TODO: HW support!
encoder = Some(
H264Encoder::new_software(size, 60, 3 * (1000 * 1000))
.expect("Failed to create encoder"),
);
encoder = Some(H264Encoder::new_nvenc(
&dev,
size.clone(),
60,
5 * (1000 * 1000),
)?);
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(

View file

@ -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:
@ -161,6 +166,9 @@ impl H264Encoder {
dict.set("forced-idr", "1");
// damn you
dict.set("delay", "0");
let encoder = video_encoder_context
.open_as_with(encoder, dict)
.with_context(|| "While opening h264_nvenc video codec")?;
@ -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.");
}
};
}