enable nvenc encoding

graynk · Jul 21, 2024 · 637cf17 · 637cf17
1 parent 37fd1a9
commit 637cf17
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 10 deletions.
diff --git a/app/distorters/animation.go b/app/distorters/animation.go
@@ -23,7 +23,7 @@ const (
 	Queued  = "Your message has been queued"
 )
 
-func DistortVideo(filename, output string, progressChan chan string) {
+func DistortVideo(filename, codec, output string, progressChan chan string) {
 	progressChan <- "Extracting frames..."
 	defer close(progressChan)
 	framesDir := filename + "Frames"
@@ -68,7 +68,7 @@ func DistortVideo(filename, output string, progressChan chan string) {
 		}
 	}
 	progressChan <- "Collecting frames..."
-	err = collectFramesToVideo(numberedFileName, frameRateFraction, output)
+	err = collectFramesToVideo(numberedFileName, frameRateFraction, codec, output)
 	if err != nil {
 		progressChan <- Failed
 	}
@@ -116,11 +116,11 @@ func extractFramesFromVideoSticker(frameRateFraction, filename, numberedFileName
 		numberedFileName)
 }
 
-func collectFramesToVideo(numberedFileName, frameRateFraction, filename string) error {
+func collectFramesToVideo(numberedFileName, frameRateFraction, codec, filename string) error {
 	return runFfmpeg("-r", frameRateFraction,
 		"-i", numberedFileName,
 		"-f", "mp4",
-		"-c:v", "libx264",
+		"-c:v", codec,
 		"-an",
 		"-pix_fmt", "yuv420p",
 		filename)

diff --git a/app/distorters/image.go b/app/distorters/image.go
@@ -10,10 +10,11 @@ import (
 
 func DistortImage(path string) error {
 	cmd := exec.Command(
-		"mogrify",
-		"-scale", "512x512>", // A reasonable cutoff, I hope
+		"magick",
+		path,
+		"-resize", "512x512>", // A reasonable cutoff, I hope
 		"-liquid-rescale", "50%",
-		"-scale", "200%",
+		"-resize", "200%",
 		path)
 	cmd.SysProcAttr = &syscall.SysProcAttr{
 		Setpgid: true,

diff --git a/app/distortioner.go b/app/distortioner.go
@@ -31,6 +31,7 @@ type DistorterBot struct {
 	mu          *sync.Mutex
 	graceWg     *sync.WaitGroup
 	videoWorker *tools.VideoWorker
+	codec       string
 }
 
 func (d DistorterBot) handleAnimationDistortion(c tb.Context) error {
@@ -328,13 +329,18 @@ func main() {
 	if err != nil {
 		logger.Fatal(err)
 	}
+	codec := os.Getenv("DISTORTIONER_CODEC")
+	if codec == "" {
+		codec = "libx264"
+	}
 	d := DistorterBot{
 		adminID:     adminID,
 		rl:          tools.NewRateLimiter(),
 		logger:      logger,
 		mu:          &sync.Mutex{},
 		graceWg:     &sync.WaitGroup{},
 		videoWorker: tools.NewVideoWorker(3),
+		codec:       codec,
 	}
 	b.Poller = tb.NewMiddlewarePoller(&tb.LongPoller{Timeout: 10 * time.Second}, func(update *tb.Update) bool {
 		if update.Message == nil {

diff --git a/app/handler_helpers.go b/app/handler_helpers.go
@@ -39,7 +39,7 @@ func (d DistorterBot) HandleAnimationCommon(c tb.Context) (*tb.Message, string,
 	}
 	animationOutput := filename + ".mp4"
 	progressChan := make(chan string, 3)
-	go distorters.DistortVideo(filename, animationOutput, progressChan)
+	go distorters.DistortVideo(filename, d.codec, animationOutput, progressChan)
 	for report := range progressChan {
 		if progressMessage == nil {
 			continue

diff --git a/app/tools/video_worker.go b/app/tools/video_worker.go
@@ -51,5 +51,5 @@ func (vw *VideoWorker) QueueStats() (int, int) {
 }
 
 func (vw *VideoWorker) IsBusy() bool {
-	return vw.queue.Len() > 0
+	return vw.queue.Len() > vw.workerCount
 }
diff --git a/distortioner.env b/distortioner.env
@@ -1,2 +1,3 @@
 DISTORTIONER_BOT_TOKEN=YOUR_BOT_TOKEN
-DISTORTIONER_ADMIN_ID=YOUR_TELEGRAM_ID
+DISTORTIONER_ADMIN_ID=YOUR_TELEGRAM_ID
+DISTORTIONER_CODEC=h264_nvenc