feature: ffmpeg cuda option

Author: alxarno

cmd/tinytune/tinytune.go

@@ -95,6 +95,13 @@ func main() {
 				Destination: &rawConfig.Video,
 				Category:    ProcessingCLICategory,
 			},
+			&cli.StringFlag{
+				Name:        "video-processing-accel",
+				Value:       string(preview.Auto),
+				Usage:       "processing type for videos: 'auto', 'hardware', 'software'",
+				Destination: &rawConfig.VideoProcessingAccel,
+				Category:    ProcessingCLICategory,
+			},
 			&cli.BoolFlag{
 				Name:        "image",
 				Value:       rawConfig.Images,
@@ -240,6 +247,7 @@ func start(config internal.Config) {
 		preview.WithMaxVideos(config.Process.Video.MaxItems),
 		preview.WithMaxFileSize(config.Process.MaxFileSize),
 		preview.WithTimeout(config.Process.Timeout),
+		preview.WithVideoAccel(config.Process.VideoAccel),
 	)
 	internal.PanicError(err)
 

internal/cli.go

@@ -10,25 +10,27 @@ import (
 	"time"
 
 	"github.com/alxarno/tinytune/pkg/bytesutil"
+	"github.com/alxarno/tinytune/pkg/preview"
 )
 
 const defaultPort = 8080
 const maxPortNumber = 65536
 
 type RawConfig struct {
-	Dir           string
-	Parallel      int
-	Video         bool
-	Images        bool
-	MaxImages     int64
-	MaxVideos     int64
-	Includes      string
-	Excludes      string
-	MaxFileSize   string
-	Streaming     string
-	MediaTimeout  string
-	IndexFileSave bool
-	Port          int
+	Dir                  string
+	Parallel             int
+	Video                bool
+	VideoProcessingAccel string
+	Images               bool
+	MaxImages            int64
+	MaxVideos            int64
+	Includes             string
+	Excludes             string
+	MaxFileSize          string
+	Streaming            string
+	MediaTimeout         string
+	IndexFileSave        bool
+	Port                 int
 }
 
 type MediaTypeConfig struct {
@@ -54,6 +56,7 @@ func (c MediaTypeConfig) Print(name string) {
 type ProcessConfig struct {
 	Parallel    int
 	Video       MediaTypeConfig
+	VideoAccel  preview.VideoProcessingAccelType
 	Timeout     time.Duration
 	Image       MediaTypeConfig
 	Includes    []*regexp.Regexp
@@ -90,7 +93,11 @@ func (c ProcessConfig) Print() {
 	}
 
 	if c.MaxFileSize != -1 {
-		slog.String("max-file-size", bytesutil.PrettyByteSize(c.MaxFileSize))
+		params = append(params, slog.String("max-file-size", bytesutil.PrettyByteSize(c.MaxFileSize)))
+	}
+
+	if c.VideoAccel != preview.Auto {
+		params = append(params, slog.String("video-processing-accel", string(c.VideoAccel)))
 	}
 
 	slog.Info(
@@ -128,8 +135,10 @@ func (c Config) Print() {
 
 func DefaultRawConfig() RawConfig {
 	return RawConfig{
-		Dir:           os.Getenv("PWD"),
-		Parallel:      runtime.NumCPU(),
+		Dir:                  os.Getenv("PWD"),
+		Parallel:             runtime.NumCPU(),
+		VideoProcessingAccel: string(preview.Auto),
+		// Parallel:             8,
 		Port:          defaultPort,
 		Video:         true,
 		Images:        true,
@@ -156,6 +165,7 @@ func NewConfig(raw RawConfig) Config {
 			Timeout:     getDuration(raw.MediaTimeout),
 			Parallel:    raw.Parallel,
 			Video:       MediaTypeConfig{raw.Video, raw.MaxVideos},
+			VideoAccel:  preview.VideoProcessingAccelType(raw.VideoProcessingAccel),
 			Image:       MediaTypeConfig{raw.Images, raw.MaxImages},
 			Includes:    getRegularExpressions(raw.Includes),
 			Excludes:    getRegularExpressions(raw.Excludes),

pkg/index/builder.go

@@ -20,7 +20,7 @@ var (
 )
 
 type PreviewGenerator interface {
-	Pull(item preview.Source) (preview.Data, error)
+	Pull(ctx context.Context, item preview.Source) (preview.Data, error)
 	Close()
 }
 
@@ -131,7 +131,7 @@ func (ib *indexBuilder) loadFile(
 		defer sem.Release(1)
 		defer wg.Done()
 
-		preview, err := ib.params.preview.Pull(metaItem)
+		preview, err := ib.params.preview.Pull(ctx, metaItem)
 		if err != nil {
 			slog.Error(fmt.Errorf("%w (%s): %w", ErrPreviewPull, metaItem.RelativePath, err).Error())
 			dst <- loadedFile{metaItem, nil}

pkg/index/index_test.go

@@ -122,7 +122,7 @@ func (m mockPreviewData) Duration() time.Duration {
 }
 
 //nolint:ireturn
-func (mock mockPreviewGenerator) Pull(_ preview.Source) (preview.Data, error) {
+func (mock mockPreviewGenerator) Pull(_ context.Context, _ preview.Source) (preview.Data, error) {
 	return mockPreviewData{data: mock.sampleData}, nil
 }
 

pkg/preview/ffmpeg.go

@@ -0,0 +1,57 @@
+package preview
+
+import (
+	"fmt"
+	"slices"
+	"time"
+
+	"github.com/alxarno/tinytune/pkg/timeutil"
+)
+
+type ffmpegHWAccelType string
+
+const (
+	ffmpegSoftwareAccel ffmpegHWAccelType = "software"
+	ffmpegCudaAccel     ffmpegHWAccelType = "cuda"
+)
+
+func cudaOptions(path string, start time.Duration, index int) []string {
+	mapValue := fmt.Sprintf("%d:v:0", index)
+	cudaOptions := []string{"-hwaccel", "cuda", "-hwaccel_output_format", "cuda"}
+	inputOptions := []string{"-ss", timeutil.String(start), "-i", path}
+	filterOptions := []string{"-vf", "scale_cuda=256:-1", "-frames:v", "1"}
+	encoderOptions := []string{"-c:v", "h264_nvenc"}
+	outputOptions := []string{"-an", "-f", "rawvideo", "-map", mapValue, "pipe:1"}
+
+	return slices.Concat(cudaOptions, inputOptions, filterOptions, encoderOptions, outputOptions)
+}
+
+func swOptions(path string, start time.Duration, index int) []string {
+	mapValue := fmt.Sprintf("%d:v:0", index)
+	inputOptions := []string{"-ss", timeutil.String(start), "-i", path}
+	filterOptions := []string{"-vf", "scale=256:-2", "-frames:v", "1"}
+	encoderOptions := []string{"-c:v", "libx264", "-preset", "ultrafast", "-tune", "zerolatency", "-crf", "0"}
+	outputOptions := []string{"-an", "-f", "rawvideo", "-map", mapValue, "pipe:1"}
+
+	return slices.Concat(inputOptions, filterOptions, encoderOptions, outputOptions)
+}
+
+func options(accel ffmpegHWAccelType) func(path string, start time.Duration, index int) []string {
+	switch accel {
+	case ffmpegCudaAccel:
+		return cudaOptions
+	case ffmpegSoftwareAccel:
+		return swOptions
+	default:
+		return swOptions
+	}
+}
+
+func tileOptions() []string {
+	quietOptions := []string{"-hide_banner", "-loglevel", "error"}
+	inputOptions := []string{"-i", "pipe:0"}
+	filterOptions := []string{"-c:v", "libwebp", "-vf", "tile=1x5", "-frames:v", "1"}
+	outputOptions := []string{"-f", "image2", "-an", "pipe:1"}
+
+	return slices.Concat(quietOptions, inputOptions, filterOptions, outputOptions)
+}

pkg/preview/ffmpeg_prob.go

@@ -0,0 +1,197 @@
+package preview
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os/exec"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/hashicorp/go-version"
+)
+
+var (
+	ErrCommandNotFound          = errors.New("not found")
+	ErrIncorrectFFmpegVersion   = errors.New("can't be parsed")
+	ErrOutdatedFFmpegVersion    = errors.New("is outdated")
+	ErrMetaInfoUnmarshal        = errors.New("failed to decode file's meta information")
+	ErrMetaInfoFramesCountParse = errors.New("failed to parse frames count from meta information")
+	ErrMetaInfoDurationParse    = errors.New("failed to parse duration from meta information")
+	ErrVideoStreamNotFound      = errors.New("video stream not found")
+	ErrParseFrameRate           = errors.New("failed to parse frame rate")
+	ErrNoSupportedCodecs        = errors.New("there are no supported codecs")
+)
+
+type probeFormat struct {
+	Duration string `json:"duration"`
+}
+
+type probeStream struct {
+	Frames       string `json:"nb_frames"` //nolint:tagliatelle
+	Width        int    `json:"width"`
+	Height       int    `json:"height"`
+	AvgFrameRate string `json:"avg_frame_rate"` //nolint:tagliatelle
+	CodecType    string `json:"codec_type"`     //nolint:tagliatelle
+	CodecName    string `json:"codec_nae"`      //nolint:tagliatelle
+}
+
+type probeData struct {
+	Format  probeFormat   `json:"format"`
+	Streams []probeStream `json:"streams"`
+}
+
+func processorProbe() error {
+	ctx := context.Background()
+
+	if err := probeFFmpeg(ctx, "ffmpeg"); err != nil {
+		return err
+	}
+
+	if err := probeFFmpeg(ctx, "ffprobe"); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func formatCodecs(buff *bytes.Buffer) []string {
+	codecs := []string{}
+	scanner := bufio.NewScanner(buff)
+	scanner.Split(bufio.ScanLines)
+
+	for scanner.Scan() {
+		withoutPrefix := strings.TrimPrefix(scanner.Text(), "(codec ")
+		codec := strings.TrimSuffix(withoutPrefix, ")")
+		codecs = append(codecs, codec)
+	}
+
+	return codecs
+}
+
+func probeCuda(ctx context.Context) ([]string, error) {
+	//nolint:lll
+	cmd := exec.CommandContext(ctx, "bash", "-c", "ffprobe -hide_banner -decoders | grep Nvidia | grep -oP \"[(]codec\\s\\w+[)]\"")
+	outBuff := bytes.NewBuffer(nil)
+	errBuff := bytes.NewBuffer(nil)
+	cmd.Stdout = outBuff
+	cmd.Stderr = errBuff
+
+	if err := cmd.Run(); err != nil {
+		return nil, fmt.Errorf("[%s] %w", errBuff.String(), err)
+	}
+
+	if outBuff.Len() == 0 {
+		return nil, fmt.Errorf("'%s' %w", "ffprobe -hide_banner -decoders", ErrCommandNotFound)
+	}
+
+	supportedCudaCodecs := formatCodecs(outBuff)
+	if len(supportedCudaCodecs) == 0 {
+		return nil, ErrNoSupportedCodecs
+	}
+
+	return supportedCudaCodecs, nil
+}
+
+func probeFFmpeg(ctx context.Context, com string) error {
+	cmd := exec.CommandContext(ctx, "bash", "-c", com+" -version | sed -n \"s/"+com+" version \\([-0-9.]*\\).*/\\1/p;\"") //nolint:gosec,lll
+	buf := bytes.NewBuffer(nil)
+	stdErrBuf := bytes.NewBuffer(nil)
+	cmd.Stdout = buf
+	cmd.Stderr = stdErrBuf
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("[%s] %w", stdErrBuf.String(), err)
+	}
+
+	if buf.Len() == 0 {
+		return fmt.Errorf("'%s' %w", com, ErrCommandNotFound)
+	}
+
+	clearVersion := strings.TrimSuffix(strings.TrimSuffix(buf.String(), "\n"), "-0")
+	required, _ := version.NewVersion("4.4.2")
+
+	existed, err := version.NewVersion(clearVersion)
+	if err != nil {
+		return fmt.Errorf("%w version(%s) of %s: %w", ErrIncorrectFFmpegVersion, buf.String(), com, err)
+	}
+
+	if existed.LessThan(required) {
+		return fmt.Errorf("%w version(%s) of %s: %w", ErrOutdatedFFmpegVersion, existed.String(), com, err)
+	}
+
+	return nil
+}
+
+func getVideoStream(streams []probeStream) *probeStream {
+	for _, v := range streams {
+		if v.CodecType == "video" {
+			return &v
+		}
+	}
+
+	return nil
+}
+
+type probeOutput struct {
+	width    int
+	height   int
+	duration time.Duration
+	codec    string
+}
+
+func probeOutputFrames(a string) (probeOutput, error) {
+	data := probeData{}
+	output := probeOutput{}
+
+	if err := json.Unmarshal([]byte(a), &data); err != nil {
+		return output, fmt.Errorf("%w: %w", ErrMetaInfoUnmarshal, err)
+	}
+
+	seconds, err := strconv.ParseFloat(data.Format.Duration, 64)
+	if err != nil {
+		return output, fmt.Errorf("%w: %w", ErrMetaInfoDurationParse, err)
+	}
+
+	videoStream := getVideoStream(data.Streams)
+	if videoStream == nil {
+		return output, ErrVideoStreamNotFound
+	}
+
+	output.width = videoStream.Width
+	output.height = videoStream.Height
+	output.duration = time.Duration(seconds) * time.Second
+	output.codec = videoStream.CodecName
+
+	return output, nil
+}
+
+func videoProbe(ctx context.Context, path string, timeOut time.Duration) (string, error) {
+	if timeOut > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, timeOut)
+		defer cancel()
+	}
+
+	logOptions := []string{"-hide_banner", "-loglevel", "quiet"}
+	jobOptions := []string{"-show_format", "-show_streams", "-of", "json", path}
+	options := []string{}
+	options = append(options, logOptions...)
+	options = append(options, jobOptions...)
+
+	cmd := exec.CommandContext(ctx, "ffprobe", options...)
+	buf := bytes.NewBuffer(nil)
+	stdErrBuf := bytes.NewBuffer(nil)
+	cmd.Stdout = buf
+	cmd.Stderr = stdErrBuf
+
+	if err := cmd.Run(); err != nil {
+		return "", fmt.Errorf("[%s] %w", stdErrBuf.String(), err)
+	}
+
+	return buf.String(), nil
+}