Proper check if nvenc is supported

1 files changed, 45 insertions, 21 deletions

diff --git a/src/main.cpp b/src/main.cpp
index 28c4079..a405023 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -556,7 +556,7 @@ static void receive_frames(AVCodecContext *av_codec_context, int stream_index, A
     }
 }
 
-static AVCodecContext* create_audio_codec_context(AVFormatContext *av_format_context, int fps) {
+static AVCodecContext* create_audio_codec_context(int fps) {
     const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_AAC);
     if (!codec) {
         fprintf(
@@ -590,27 +590,12 @@ static AVCodecContext* create_audio_codec_context(AVFormatContext *av_format_con
     codec_context->framerate.num = fps;
     codec_context->framerate.den = 1;
 
-    av_format_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
     codec_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
 
     return codec_context;
 }
 
-static const AVCodec* find_h264_encoder() {
-    const AVCodec *codec = avcodec_find_encoder_by_name("h264_nvenc");
-    if(!codec)
-        codec = avcodec_find_encoder_by_name("nvenc_h264");
-    return codec;
-}
-
-static const AVCodec* find_h265_encoder() {
-    const AVCodec *codec = avcodec_find_encoder_by_name("hevc_nvenc");
-    if(!codec)
-        codec = avcodec_find_encoder_by_name("nvenc_hevc");
-    return codec;
-}
-
-static AVCodecContext *create_video_codec_context(AVFormatContext *av_format_context, 
+static AVCodecContext *create_video_codec_context(AVPixelFormat pix_fmt,
                             VideoQuality video_quality,
                             int record_width, int record_height,
                             int fps, const AVCodec *codec, bool is_livestream) {
@@ -644,7 +629,7 @@ static AVCodecContext *create_video_codec_context(AVFormatContext *av_format_con
         codec_context->gop_size = fps * 2;
     }
     codec_context->max_b_frames = 0;
-    codec_context->pix_fmt = AV_PIX_FMT_CUDA;
+    codec_context->pix_fmt = pix_fmt;
     codec_context->color_range = AVCOL_RANGE_JPEG;
     if(codec->id == AV_CODEC_ID_HEVC)
         codec_context->codec_tag = MKTAG('h', 'v', 'c', '1');
@@ -697,12 +682,50 @@ static AVCodecContext *create_video_codec_context(AVFormatContext *av_format_con
     //codec_context->rc_min_rate = codec_context->bit_rate;
     //codec_context->rc_buffer_size = codec_context->bit_rate / 10;
 
-    av_format_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
     codec_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
 
     return codec_context;
 }
 
+static const AVCodec* find_h264_encoder() {
+    const AVCodec *codec = avcodec_find_encoder_by_name("h264_nvenc");
+    if(!codec)
+        codec = avcodec_find_encoder_by_name("nvenc_h264");
+
+    static bool checked = false;
+    if(!checked) {
+        checked = true;
+        // Do not use AV_PIX_FMT_CUDA because we dont want to do full check with hardware context
+        AVCodecContext *codec_context = create_video_codec_context(AV_PIX_FMT_YUV420P, VideoQuality::VERY_HIGH, 1920, 1080, 60, codec, false);
+        if (avcodec_open2(codec_context, codec_context->codec, NULL) < 0) {
+            avcodec_free_context(&codec_context);
+            return nullptr;
+        }
+    }
+    return codec;
+}
+
+static const AVCodec* find_h265_encoder() {
+    const AVCodec *codec = avcodec_find_encoder_by_name("hevc_nvenc");
+    if(!codec)
+        codec = avcodec_find_encoder_by_name("nvenc_hevc");
+
+    if(!codec)
+        return nullptr;
+
+    static bool checked = false;
+    if(!checked) {
+        checked = true;
+        // Do not use AV_PIX_FMT_CUDA because we dont want to do full check with hardware context
+        AVCodecContext *codec_context = create_video_codec_context(AV_PIX_FMT_YUV420P, VideoQuality::VERY_HIGH, 1920, 1080, 60, codec, false);
+        if (avcodec_open2(codec_context, codec_context->codec, NULL) < 0) {
+            avcodec_free_context(&codec_context);
+            return nullptr;
+        }
+    }
+    return codec;
+}
+
 static AVFrame* open_audio(AVCodecContext *audio_codec_context) {
     int ret;
     ret = avcodec_open2(audio_codec_context, audio_codec_context->codec, nullptr);
@@ -1522,6 +1545,7 @@ int main(int argc, char **argv) {
         return 1;
     }
 
+    av_format_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
     av_format_context->flags |= AVFMT_FLAG_GENPTS;
     const AVOutputFormat *output_format = av_format_context->oformat;
 
@@ -1536,7 +1560,7 @@ int main(int argc, char **argv) {
     AVStream *video_stream = nullptr;
     std::vector<AudioTrack> audio_tracks;
 
-    AVCodecContext *video_codec_context = create_video_codec_context(av_format_context, quality, record_width, record_height, fps, video_codec_f, is_livestream);
+    AVCodecContext *video_codec_context = create_video_codec_context(AV_PIX_FMT_CUDA, quality, record_width, record_height, fps, video_codec_f, is_livestream);
     if(replay_buffer_size_secs == -1)
         video_stream = create_stream(av_format_context, video_codec_context);
 
@@ -1548,7 +1572,7 @@ int main(int argc, char **argv) {
 
     int audio_stream_index = VIDEO_STREAM_INDEX + 1;
     for(const AudioInput &audio_input : requested_audio_inputs) {
-        AVCodecContext *audio_codec_context = create_audio_codec_context(av_format_context, fps);
+        AVCodecContext *audio_codec_context = create_audio_codec_context(fps);
 
         AVStream *audio_stream = nullptr;
         if(replay_buffer_size_secs == -1)