From 70325e02770a36a7697d33573e49f503396ce1bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Fri, 31 Mar 2023 22:17:46 +0200
Subject: [PATCH] add nvenc pipeline.

---
 docs/getting-started/configuration.md |  5 +++--
 server/internal/capture/pipelines.go  | 12 ++++++++---
 server/internal/config/capture.go     | 31 +++++++++++++++++++++------
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md
index d61b87a..3688a21 100644
--- a/docs/getting-started/configuration.md
+++ b/docs/getting-started/configuration.md
@@ -84,8 +84,9 @@ nat1to1: <ip>
   - The resulting stream frames per seconds should be capped *(0 for uncapped)*.
   - e.g. `0`
 #### `NEKO_HWENC`:
-  - Use hardware accelerated encoding, for now supported only `VAAPI`.
-  - e.g. `VAAPI`
+  - none *(default CPU encoding)*
+  - vaapi
+  - nvenc
 
 ### Audio
 
diff --git a/server/internal/capture/pipelines.go b/server/internal/capture/pipelines.go
index 4253fc4..8c6e853 100644
--- a/server/internal/capture/pipelines.go
+++ b/server/internal/capture/pipelines.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"m1k1o/neko/internal/capture/gst"
+	"m1k1o/neko/internal/config"
 	"m1k1o/neko/internal/types/codec"
 )
 
@@ -52,7 +53,7 @@ func NewBroadcastPipeline(device string, display string, pipelineSrc string, url
 	return pipelineStr, nil
 }
 
-func NewVideoPipeline(rtpCodec codec.RTPCodec, display string, pipelineSrc string, fps int16, bitrate uint, hwenc string) (string, error) {
+func NewVideoPipeline(rtpCodec codec.RTPCodec, display string, pipelineSrc string, fps int16, bitrate uint, hwenc config.HwEnc) (string, error) {
 	pipelineStr := " ! appsink name=appsinkvideo"
 
 	// if using custom pipeline
@@ -68,7 +69,7 @@ func NewVideoPipeline(rtpCodec codec.RTPCodec, display string, pipelineSrc strin
 
 	switch rtpCodec.Name {
 	case codec.VP8().Name:
-		if hwenc == "VAAPI" {
+		if hwenc == config.HwEncVAAPI {
 			if err := gst.CheckPlugins([]string{"ximagesrc", "vaapi"}); err != nil {
 				return "", err
 			}
@@ -138,13 +139,18 @@ func NewVideoPipeline(rtpCodec codec.RTPCodec, display string, pipelineSrc strin
 			return "", err
 		}
 
-		if hwenc == "VAAPI" {
+		if hwenc == config.HwEncVAAPI {
 			if err := gst.CheckPlugins([]string{"vaapi"}); err != nil {
 				return "", err
 			}
 
 			pipelineStr = fmt.Sprintf(videoSrc+"video/x-raw,format=NV12 ! vaapih264enc rate-control=vbr bitrate=%d keyframe-period=180 quality-level=7 ! video/x-h264,stream-format=byte-stream,profile=constrained-baseline"+pipelineStr, display, fps, bitrate)
+		} else if hwenc == config.HwEncNVENC {
+			if err := gst.CheckPlugins([]string{"nvcodec"}); err != nil {
+				return "", err
+			}
 
+			pipelineStr = fmt.Sprintf(videoSrc+"video/x-raw,format=NV12 ! nvh264enc name=encoder preset=5 zerolatency=1 gop-size=15 bitrate=%d rc-mode=5 ! video/x-h264,stream-format=byte-stream,profile=constrained-baseline"+pipelineStr, display, fps, bitrate)
 		} else {
 			// https://gstreamer.freedesktop.org/documentation/openh264/openh264enc.html?gi-language=c#openh264enc
 			// gstreamer1.0-plugins-bad
diff --git a/server/internal/config/capture.go b/server/internal/config/capture.go
index efb8b7d..7c4dbec 100644
--- a/server/internal/config/capture.go
+++ b/server/internal/config/capture.go
@@ -2,6 +2,7 @@ package config
 
 import (
 	"m1k1o/neko/internal/types/codec"
+	"strings"
 
 	"github.com/pion/webrtc/v3"
 	"github.com/rs/zerolog/log"
@@ -9,13 +10,21 @@ import (
 	"github.com/spf13/viper"
 )
 
+type HwEnc int
+
+const (
+	HwEncNone HwEnc = iota
+	HwEncVAAPI
+	HwEncNVENC
+)
+
 type Capture struct {
 	// video
 	Display       string
 	VideoCodec    codec.RTPCodec
-	VideoHWEnc    string // TODO: Pipeline builder.
-	VideoBitrate  uint   // TODO: Pipeline builder.
-	VideoMaxFPS   int16  // TODO: Pipeline builder.
+	VideoHWEnc    HwEnc // TODO: Pipeline builder.
+	VideoBitrate  uint  // TODO: Pipeline builder.
+	VideoMaxFPS   int16 // TODO: Pipeline builder.
 	VideoPipeline string
 
 	// audio
@@ -184,11 +193,19 @@ func (s *Capture) Set() {
 		log.Warn().Msg("you are using deprecated config setting 'NEKO_AV1=true', use 'NEKO_VIDEO_CODEC=av1' instead")
 	}
 
-	videoHWEnc := ""
-	if viper.GetString("hwenc") == "VAAPI" {
-		videoHWEnc = "VAAPI"
+	videoHWEnc := strings.ToLower(viper.GetString("hwenc"))
+	switch videoHWEnc {
+	case "":
+		fallthrough
+	case "none":
+		s.VideoHWEnc = HwEncNone
+	case "vaapi":
+		s.VideoHWEnc = HwEncVAAPI
+	case "nvenc":
+		s.VideoHWEnc = HwEncNVENC
+	default:
+		log.Warn().Str("hwenc", videoHWEnc).Msgf("unknown video hw encoder, using CPU")
 	}
-	s.VideoHWEnc = videoHWEnc
 
 	s.VideoBitrate = viper.GetUint("video_bitrate")
 	s.VideoMaxFPS = int16(viper.GetInt("max_fps"))