Add experimental av1 support (only tested on amd)

2 files changed, 71 insertions, 12 deletions

diff --git a/TODO b/TODO
index 3e610f0..bf8a413 100644
--- a/TODO
+++ b/TODO
@@ -54,7 +54,7 @@ Intel is a bit weird with monitor capture and multiple monitors. If one of the m
     Is that only the case when the primary monitor is rotated? Also the primary monitor becomes position 0, 0 so crtc (x11 randr) position doesn't match the drm pos. Maybe get monitor position and size from drm instead.
     How about if multiple monitors are rotated?
 
-Support vp8/vp9/av1. This is especially important on amd which on some distros (such as Manjaro) where hardware accelerated h264/hevc is disabled in the mesa package.
+Support vp8/vp9. This is especially important on amd which on some distros (such as Manjaro) where hardware accelerated h264/hevc is disabled in the mesa package.
 
 Support screen (all monitors) capture on amd/intel when no combined plane is found.
 Use separate plane (which has offset and pitch) from combined plane instead of the combined plane.
@@ -96,8 +96,6 @@ Remove follow focused option.
 
 Fix kms capture on 10-bit setup (both x11 and wayland). It may already work automatically on amd/intel, but on nvidia the color conversion shader needs to be used to automatically map 10-bit to 8-bit.
 
-Test vaapi av1 encoding (av1_vaapi) (which was recently added to ffmpeg master).
-
 Overclocking (-oc) can overclock too much on some systems. Maybe remove the option?
 
 Exit if X11/Wayland killed (if drm plane dead or something?)
@@ -107,3 +105,5 @@ Use SRC_W and SRC_H for screen plane instead of crtc_w and crtc_h.
 Make it possible to select which /dev/dri/card* to use, but that requires opengl to also use the same card. Not sure if that is possible for amd, intel and nvidia without using vulkan instead.
 
 Support I915_FORMAT_MOD_Y_TILED_CCS (and other power saving modifiers, see https://trac.ffmpeg.org/ticket/8542). The only fix may be to use desktop portal for recording. This issue doesn't appear on x11 since these modifiers are not used by xorg server.
+
+Test if p2 state can be worked around by using pure nvenc api and overwriting cuInit/cuCtxCreate* to not do anything. Cuda might be loaded when using nvenc but it might not be used, with certain record options? (such as h264 p5).
\ No newline at end of file
diff --git a/src/main.cpp b/src/main.cpp
index 53d936c..457728e 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -67,7 +67,8 @@ enum class VideoQuality {
 
 enum class VideoCodec {
     H264,
-    H265
+    H265,
+    AV1
 };
 
 enum class AudioCodec {
@@ -501,6 +502,24 @@ static const AVCodec* find_h265_encoder(gsr_gpu_vendor vendor, const char *card_
     return checked_success ? codec : nullptr;
 }
 
+static const AVCodec* find_av1_encoder(gsr_gpu_vendor vendor, const char *card_path) {
+    const AVCodec *codec = avcodec_find_encoder_by_name(vendor == GSR_GPU_VENDOR_NVIDIA ? "av1_nvenc" : "av1_vaapi");
+    if(!codec)
+        codec = avcodec_find_encoder_by_name(vendor == GSR_GPU_VENDOR_NVIDIA ? "nvenc_av1" : "vaapi_av1");
+
+    if(!codec)
+        return nullptr;
+
+    static bool checked = false;
+    static bool checked_success = true;
+    if(!checked) {
+        checked = true;
+        if(!check_if_codec_valid_for_hardware(codec, vendor, card_path))
+            checked_success = false;
+    }
+    return checked_success ? codec : nullptr;
+}
+
 static AVFrame* open_audio(AVCodecContext *audio_codec_context) {
     AVDictionary *options = nullptr;
     av_dict_set(&options, "strict", "experimental", 0);
@@ -620,6 +639,15 @@ static void open_video(AVCodecContext *codec_context, VideoQuality video_quality
                     av_dict_set(&options, "profile", "high444p", 0);
                     break;
             }
+        } else if(codec_context->codec_id == AV_CODEC_ID_AV1) {
+            switch(pixel_format) {
+                case PixelFormat::YUV420:
+                    av_dict_set(&options, "rgb_mode", "yuv420", 0);
+                    break;
+                case PixelFormat::YUV444:
+                    av_dict_set(&options, "rgb_mode", "yuv444", 0);
+                    break;
+            }
         } else {
             //av_dict_set(&options, "profile", "main10", 0);
             //av_dict_set(&options, "pix_fmt", "yuv420p16le", 0);
@@ -647,6 +675,10 @@ static void open_video(AVCodecContext *codec_context, VideoQuality video_quality
         if(codec_context->codec_id == AV_CODEC_ID_H264) {
             av_dict_set(&options, "profile", "high", 0);
             av_dict_set_int(&options, "quality", 7, 0);
+        } else if(codec_context->codec_id == AV_CODEC_ID_AV1) {
+            av_dict_set(&options, "profile", "main", 0); // TODO: use professional instead?
+            av_dict_set(&options, "tier", "main", 0);
+            av_dict_set_int(&options, "quality", 7, 0);
         } else {
             av_dict_set(&options, "profile", "main", 0);
         }
@@ -666,7 +698,7 @@ static void open_video(AVCodecContext *codec_context, VideoQuality video_quality
 }
 
 static void usage_header() {
-    fprintf(stderr, "usage: gpu-screen-recorder -w <window_id|monitor|focused> [-c <container_format>] [-s WxH] -f <fps> [-a <audio_input>] [-q <quality>] [-r <replay_buffer_size_sec>] [-k h264|h265] [-ac aac|opus|flac] [-oc yes|no] [-fm cfr|vfr] [-v yes|no] [-h|--help] [-o <output_file>] [-mf yes|no]\n");
+    fprintf(stderr, "usage: gpu-screen-recorder -w <window_id|monitor|focused> [-c <container_format>] [-s WxH] -f <fps> [-a <audio_input>] [-q <quality>] [-r <replay_buffer_size_sec>] [-k h264|h265|av1] [-ac aac|opus|flac] [-oc yes|no] [-fm cfr|vfr] [-v yes|no] [-h|--help] [-o <output_file>] [-mf yes|no]\n");
 }
 
 static void usage_full() {
@@ -683,7 +715,7 @@ static void usage_full() {
     fprintf(stderr, "\n");
     fprintf(stderr, "  -c    Container format for output file, for example mp4, or flv. Only required if no output file is specified or if recording in replay buffer mode.\n");
     fprintf(stderr, "        If an output file is specified and -c is not used then the container format is determined from the output filename extension.\n");
-    fprintf(stderr, "        Only containers that support h264 or hevc are supported, which means that only mp4, mkv, flv (and some others) are supported.\n");
+    fprintf(stderr, "        Only containers that support h264, hevc or av1 are supported, which means that only mp4, mkv, flv (and some others) are supported.\n");
     fprintf(stderr, "        WebM is not supported yet.\n");
     fprintf(stderr, "\n");
     fprintf(stderr, "  -s    The size (area) to record at in the format WxH, for example 1920x1080. This option is only supported (and required) when -w is \"focused\".\n");
@@ -702,7 +734,7 @@ static void usage_full() {
     fprintf(stderr, "        and the video will only be saved when the gpu-screen-recorder is closed. This feature is similar to Nvidia's instant replay feature.\n");
     fprintf(stderr, "        This option has be between 5 and 1200. Note that the replay buffer size will not always be precise, because of keyframes. Optional, disabled by default.\n");
     fprintf(stderr, "\n");
-    fprintf(stderr, "  -k    Video codec to use. Should be either 'auto', 'h264' or 'h265'. Defaults to 'auto' which defaults to 'h265' unless recording at fps higher than 60. Defaults to 'h264' on intel.\n");
+    fprintf(stderr, "  -k    Video codec to use. Should be either 'auto', 'h264', 'h265', 'av1'. Defaults to 'auto' which defaults to 'h265' unless recording at fps higher than 60. Defaults to 'h264' on intel.\n");
     fprintf(stderr, "        Forcefully set to 'h264' if -c is 'flv'.\n");
     fprintf(stderr, "\n");
     fprintf(stderr, "  -ac   Audio codec to use. Should be either 'aac', 'opus' or 'flac'. Defaults to 'opus' for .mp4/.mkv files, otherwise defaults to 'aac'.\n");
@@ -1248,8 +1280,10 @@ int main(int argc, char **argv) {
         video_codec = VideoCodec::H264;
     } else if(strcmp(video_codec_to_use, "h265") == 0) {
         video_codec = VideoCodec::H265;
+    } else if(strcmp(video_codec_to_use, "av1") == 0) {
+        video_codec = VideoCodec::AV1;
     } else if(strcmp(video_codec_to_use, "auto") != 0) {
-        fprintf(stderr, "Error: -k should either be either 'auto', 'h264' or 'h265', got: '%s'\n", video_codec_to_use);
+        fprintf(stderr, "Error: -k should either be either 'auto', 'h264', 'h265' or 'av1', got: '%s'\n", video_codec_to_use);
         usage();
     }
 
@@ -1753,7 +1787,7 @@ int main(int argc, char **argv) {
         }
     }
 
-    //bool use_hevc = strcmp(window_str, "screen") == 0 || strcmp(window_str, "screen-direct") == 0;
+    // TODO: Allow hevc, vp9 and av1 in (enhanced) flv (supported since ffmpeg 6.1)
     const bool is_flv = strcmp(file_extension.c_str(), "flv") == 0;
     if(video_codec != VideoCodec::H264 && is_flv) {
         video_codec_to_use = "h264";
@@ -1769,6 +1803,9 @@ int main(int argc, char **argv) {
         case VideoCodec::H265:
             video_codec_f = find_h265_encoder(gpu_inf.vendor, card_path);
             break;
+        case VideoCodec::AV1:
+            video_codec_f = find_av1_encoder(gpu_inf.vendor, card_path);
+            break;
     }
 
     if(!video_codec_auto && !video_codec_f && !is_flv) {
@@ -1787,16 +1824,38 @@ int main(int argc, char **argv) {
                 video_codec_f = find_h264_encoder(gpu_inf.vendor, card_path);
                 break;
             }
+            case VideoCodec::AV1: {
+                fprintf(stderr, "Warning: selected video codec av1 is not supported, trying h264 instead\n");
+                video_codec_to_use = "h264";
+                video_codec = VideoCodec::H264;
+                video_codec_f = find_h264_encoder(gpu_inf.vendor, card_path);
+                break;
+            }
         }
     }
 
     if(!video_codec_f) {
-        const char *video_codec_name = video_codec == VideoCodec::H264 ? "h264" : "h265";
+        const char *video_codec_name = "";
+        switch(video_codec) {
+            case VideoCodec::H264: {
+                video_codec_name = "h265";
+                break;
+            }
+            case VideoCodec::H265: {
+                video_codec_name = "h265";
+                break;
+            }
+            case VideoCodec::AV1: {
+                video_codec_name = "av1";
+                break;
+            }
+        }
+
         fprintf(stderr, "Error: your gpu does not support '%s' video codec. If you are sure that your gpu does support '%s' video encoding and you are using an AMD/Intel GPU,\n"
             "  then it's possible that your distro has disabled hardware accelerated video encoding for '%s' video codec.\n"
             "  This may be the case on corporate distros such as Manjaro.\n"
-            "  You can test this by running 'vainfo | grep VAEntrypointEncSlice' to see if it matches any H264/HEVC profile. vainfo is part of libva-utils.\n"
-            "  On such distros, you need to manually install mesa from source to enable H264/HEVC hardware acceleration, or use a more user friendly distro.\n", video_codec_name, video_codec_name, video_codec_name);
+            "  You can test this by running 'vainfo | grep VAEntrypointEncSlice' to see if it matches any H264/HEVC/AV1 profile. vainfo is part of libva-utils.\n"
+            "  On such distros, you need to manually install mesa from source to enable H264/HEVC/AV1 hardware acceleration, or use a more user friendly distro.\n", video_codec_name, video_codec_name, video_codec_name);
         _exit(2);
     }