mirror of https://github.com/stashapp/stash.git
525 lines
13 KiB
Go
525 lines
13 KiB
Go
package manager
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/remeh/sizedwaitgroup"
|
|
"github.com/stashapp/stash/internal/manager/config"
|
|
"github.com/stashapp/stash/pkg/image"
|
|
"github.com/stashapp/stash/pkg/job"
|
|
"github.com/stashapp/stash/pkg/logger"
|
|
"github.com/stashapp/stash/pkg/models"
|
|
"github.com/stashapp/stash/pkg/scene"
|
|
"github.com/stashapp/stash/pkg/scene/generate"
|
|
"github.com/stashapp/stash/pkg/sliceutil/stringslice"
|
|
)
|
|
|
|
type GenerateMetadataInput struct {
|
|
Covers bool `json:"covers"`
|
|
Sprites bool `json:"sprites"`
|
|
Previews bool `json:"previews"`
|
|
ImagePreviews bool `json:"imagePreviews"`
|
|
PreviewOptions *GeneratePreviewOptionsInput `json:"previewOptions"`
|
|
Markers bool `json:"markers"`
|
|
MarkerImagePreviews bool `json:"markerImagePreviews"`
|
|
MarkerScreenshots bool `json:"markerScreenshots"`
|
|
Transcodes bool `json:"transcodes"`
|
|
// Generate transcodes even if not required
|
|
ForceTranscodes bool `json:"forceTranscodes"`
|
|
Phashes bool `json:"phashes"`
|
|
InteractiveHeatmapsSpeeds bool `json:"interactiveHeatmapsSpeeds"`
|
|
ClipPreviews bool `json:"clipPreviews"`
|
|
ImageThumbnails bool `json:"imageThumbnails"`
|
|
// scene ids to generate for
|
|
SceneIDs []string `json:"sceneIDs"`
|
|
// marker ids to generate for
|
|
MarkerIDs []string `json:"markerIDs"`
|
|
// overwrite existing media
|
|
Overwrite bool `json:"overwrite"`
|
|
}
|
|
|
|
type GeneratePreviewOptionsInput struct {
|
|
// Number of segments in a preview file
|
|
PreviewSegments *int `json:"previewSegments"`
|
|
// Preview segment duration, in seconds
|
|
PreviewSegmentDuration *float64 `json:"previewSegmentDuration"`
|
|
// Duration of start of video to exclude when generating previews
|
|
PreviewExcludeStart *string `json:"previewExcludeStart"`
|
|
// Duration of end of video to exclude when generating previews
|
|
PreviewExcludeEnd *string `json:"previewExcludeEnd"`
|
|
// Preset when generating preview
|
|
PreviewPreset *models.PreviewPreset `json:"previewPreset"`
|
|
}
|
|
|
|
const generateQueueSize = 200000
|
|
|
|
type GenerateJob struct {
|
|
repository models.Repository
|
|
input GenerateMetadataInput
|
|
|
|
overwrite bool
|
|
fileNamingAlgo models.HashAlgorithm
|
|
|
|
totals totalsGenerate
|
|
}
|
|
|
|
type totalsGenerate struct {
|
|
covers int64
|
|
sprites int64
|
|
previews int64
|
|
imagePreviews int64
|
|
markers int64
|
|
transcodes int64
|
|
phashes int64
|
|
interactiveHeatmapSpeeds int64
|
|
clipPreviews int64
|
|
imageThumbnails int64
|
|
|
|
tasks int
|
|
}
|
|
|
|
func (j *GenerateJob) Execute(ctx context.Context, progress *job.Progress) error {
|
|
var scenes []*models.Scene
|
|
var err error
|
|
var markers []*models.SceneMarker
|
|
|
|
j.overwrite = j.input.Overwrite
|
|
j.fileNamingAlgo = config.GetInstance().GetVideoFileNamingAlgorithm()
|
|
|
|
config := config.GetInstance()
|
|
parallelTasks := config.GetParallelTasksWithAutoDetection()
|
|
|
|
logger.Infof("Generate started with %d parallel tasks", parallelTasks)
|
|
|
|
queue := make(chan Task, generateQueueSize)
|
|
go func() {
|
|
defer close(queue)
|
|
|
|
sceneIDs, err := stringslice.StringSliceToIntSlice(j.input.SceneIDs)
|
|
if err != nil {
|
|
logger.Error(err.Error())
|
|
}
|
|
markerIDs, err := stringslice.StringSliceToIntSlice(j.input.MarkerIDs)
|
|
if err != nil {
|
|
logger.Error(err.Error())
|
|
}
|
|
|
|
g := &generate.Generator{
|
|
Encoder: instance.FFMpeg,
|
|
FFMpegConfig: instance.Config,
|
|
LockManager: instance.ReadLockManager,
|
|
MarkerPaths: instance.Paths.SceneMarkers,
|
|
ScenePaths: instance.Paths.Scene,
|
|
Overwrite: j.overwrite,
|
|
}
|
|
|
|
r := j.repository
|
|
if err := r.WithReadTxn(ctx, func(ctx context.Context) error {
|
|
qb := r.Scene
|
|
if len(j.input.SceneIDs) == 0 && len(j.input.MarkerIDs) == 0 {
|
|
j.queueTasks(ctx, g, queue)
|
|
} else {
|
|
if len(j.input.SceneIDs) > 0 {
|
|
scenes, err = qb.FindMany(ctx, sceneIDs)
|
|
for _, s := range scenes {
|
|
if err := s.LoadFiles(ctx, qb); err != nil {
|
|
return err
|
|
}
|
|
|
|
j.queueSceneJobs(ctx, g, s, queue)
|
|
}
|
|
}
|
|
|
|
if len(j.input.MarkerIDs) > 0 {
|
|
markers, err = r.SceneMarker.FindMany(ctx, markerIDs)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, m := range markers {
|
|
j.queueMarkerJob(g, m, queue)
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}); err != nil && ctx.Err() == nil {
|
|
logger.Error(err.Error())
|
|
return
|
|
}
|
|
|
|
totals := j.totals
|
|
logMsg := "Generating"
|
|
if j.input.Covers {
|
|
logMsg += fmt.Sprintf(" %d covers", totals.covers)
|
|
}
|
|
if j.input.Sprites {
|
|
logMsg += fmt.Sprintf(" %d sprites", totals.sprites)
|
|
}
|
|
if j.input.Previews {
|
|
logMsg += fmt.Sprintf(" %d previews", totals.previews)
|
|
}
|
|
if j.input.ImagePreviews {
|
|
logMsg += fmt.Sprintf(" %d image previews", totals.imagePreviews)
|
|
}
|
|
if j.input.Markers {
|
|
logMsg += fmt.Sprintf(" %d markers", totals.markers)
|
|
}
|
|
if j.input.Transcodes {
|
|
logMsg += fmt.Sprintf(" %d transcodes", totals.transcodes)
|
|
}
|
|
if j.input.Phashes {
|
|
logMsg += fmt.Sprintf(" %d phashes", totals.phashes)
|
|
}
|
|
if j.input.InteractiveHeatmapsSpeeds {
|
|
logMsg += fmt.Sprintf(" %d heatmaps & speeds", totals.interactiveHeatmapSpeeds)
|
|
}
|
|
if j.input.ClipPreviews {
|
|
logMsg += fmt.Sprintf(" %d Image Clip Previews", totals.clipPreviews)
|
|
}
|
|
if j.input.ImageThumbnails {
|
|
logMsg += fmt.Sprintf(" %d Image Thumbnails", totals.imageThumbnails)
|
|
}
|
|
if logMsg == "Generating" {
|
|
logMsg = "Nothing selected to generate"
|
|
}
|
|
logger.Infof(logMsg)
|
|
|
|
progress.SetTotal(int(totals.tasks))
|
|
}()
|
|
|
|
wg := sizedwaitgroup.New(parallelTasks)
|
|
|
|
// Start measuring how long the generate has taken. (consider moving this up)
|
|
start := time.Now()
|
|
if err = instance.Paths.Generated.EnsureTmpDir(); err != nil {
|
|
logger.Warnf("could not create temporary directory: %v", err)
|
|
}
|
|
|
|
defer func() {
|
|
if err := instance.Paths.Generated.EmptyTmpDir(); err != nil {
|
|
logger.Warnf("failure emptying temporary directory: %v", err)
|
|
}
|
|
}()
|
|
|
|
for f := range queue {
|
|
if job.IsCancelled(ctx) {
|
|
break
|
|
}
|
|
|
|
wg.Add()
|
|
// #1879 - need to make a copy of f - otherwise there is a race condition
|
|
// where f is changed when the goroutine runs
|
|
localTask := f
|
|
go progress.ExecuteTask(localTask.GetDescription(), func() {
|
|
localTask.Start(ctx)
|
|
wg.Done()
|
|
progress.Increment()
|
|
})
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
if job.IsCancelled(ctx) {
|
|
logger.Info("Stopping due to user request")
|
|
return nil
|
|
}
|
|
|
|
elapsed := time.Since(start)
|
|
logger.Info(fmt.Sprintf("Generate finished (%s)", elapsed))
|
|
return nil
|
|
}
|
|
|
|
func (j *GenerateJob) queueTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
|
|
j.totals = totalsGenerate{}
|
|
|
|
j.queueScenesTasks(ctx, g, queue)
|
|
j.queueImagesTasks(ctx, g, queue)
|
|
}
|
|
|
|
func (j *GenerateJob) queueScenesTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
|
|
const batchSize = 1000
|
|
|
|
findFilter := models.BatchFindFilter(batchSize)
|
|
|
|
r := j.repository
|
|
|
|
for more := true; more; {
|
|
if job.IsCancelled(ctx) {
|
|
return
|
|
}
|
|
|
|
scenes, err := scene.Query(ctx, r.Scene, nil, findFilter)
|
|
if err != nil {
|
|
logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
|
|
return
|
|
}
|
|
|
|
for _, ss := range scenes {
|
|
if job.IsCancelled(ctx) {
|
|
return
|
|
}
|
|
|
|
if err := ss.LoadFiles(ctx, r.Scene); err != nil {
|
|
logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
|
|
return
|
|
}
|
|
|
|
j.queueSceneJobs(ctx, g, ss, queue)
|
|
}
|
|
|
|
if len(scenes) != batchSize {
|
|
more = false
|
|
} else {
|
|
*findFilter.Page++
|
|
}
|
|
}
|
|
}
|
|
|
|
func (j *GenerateJob) queueImagesTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
|
|
const batchSize = 1000
|
|
|
|
findFilter := models.BatchFindFilter(batchSize)
|
|
|
|
r := j.repository
|
|
|
|
for more := j.input.ClipPreviews || j.input.ImageThumbnails; more; {
|
|
if job.IsCancelled(ctx) {
|
|
return
|
|
}
|
|
|
|
images, err := image.Query(ctx, r.Image, nil, findFilter)
|
|
if err != nil {
|
|
logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
|
|
return
|
|
}
|
|
|
|
for _, ss := range images {
|
|
if job.IsCancelled(ctx) {
|
|
return
|
|
}
|
|
|
|
if err := ss.LoadFiles(ctx, r.Image); err != nil {
|
|
logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
|
|
return
|
|
}
|
|
|
|
j.queueImageJob(g, ss, queue)
|
|
}
|
|
|
|
if len(images) != batchSize {
|
|
more = false
|
|
} else {
|
|
*findFilter.Page++
|
|
}
|
|
}
|
|
}
|
|
|
|
func getGeneratePreviewOptions(optionsInput GeneratePreviewOptionsInput) generate.PreviewOptions {
|
|
config := config.GetInstance()
|
|
|
|
ret := generate.PreviewOptions{
|
|
Segments: config.GetPreviewSegments(),
|
|
SegmentDuration: config.GetPreviewSegmentDuration(),
|
|
ExcludeStart: config.GetPreviewExcludeStart(),
|
|
ExcludeEnd: config.GetPreviewExcludeEnd(),
|
|
Preset: config.GetPreviewPreset().String(),
|
|
Audio: config.GetPreviewAudio(),
|
|
}
|
|
|
|
if optionsInput.PreviewSegments != nil {
|
|
ret.Segments = *optionsInput.PreviewSegments
|
|
}
|
|
|
|
if optionsInput.PreviewSegmentDuration != nil {
|
|
ret.SegmentDuration = *optionsInput.PreviewSegmentDuration
|
|
}
|
|
|
|
if optionsInput.PreviewExcludeStart != nil {
|
|
ret.ExcludeStart = *optionsInput.PreviewExcludeStart
|
|
}
|
|
|
|
if optionsInput.PreviewExcludeEnd != nil {
|
|
ret.ExcludeEnd = *optionsInput.PreviewExcludeEnd
|
|
}
|
|
|
|
if optionsInput.PreviewPreset != nil {
|
|
ret.Preset = optionsInput.PreviewPreset.String()
|
|
}
|
|
|
|
return ret
|
|
}
|
|
|
|
func (j *GenerateJob) queueSceneJobs(ctx context.Context, g *generate.Generator, scene *models.Scene, queue chan<- Task) {
|
|
r := j.repository
|
|
|
|
if j.input.Covers {
|
|
task := &GenerateCoverTask{
|
|
repository: r,
|
|
Scene: *scene,
|
|
Overwrite: j.overwrite,
|
|
}
|
|
|
|
if task.required(ctx) {
|
|
j.totals.covers++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
if j.input.Sprites {
|
|
task := &GenerateSpriteTask{
|
|
Scene: *scene,
|
|
Overwrite: j.overwrite,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
}
|
|
|
|
if task.required() {
|
|
j.totals.sprites++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
generatePreviewOptions := j.input.PreviewOptions
|
|
if generatePreviewOptions == nil {
|
|
generatePreviewOptions = &GeneratePreviewOptionsInput{}
|
|
}
|
|
options := getGeneratePreviewOptions(*generatePreviewOptions)
|
|
|
|
if j.input.Previews {
|
|
task := &GeneratePreviewTask{
|
|
Scene: *scene,
|
|
ImagePreview: j.input.ImagePreviews,
|
|
Options: options,
|
|
Overwrite: j.overwrite,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
generator: g,
|
|
}
|
|
|
|
if task.required() {
|
|
if task.videoPreviewRequired() {
|
|
j.totals.previews++
|
|
}
|
|
if task.imagePreviewRequired() {
|
|
j.totals.imagePreviews++
|
|
}
|
|
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
if j.input.Markers {
|
|
task := &GenerateMarkersTask{
|
|
repository: r,
|
|
Scene: scene,
|
|
Overwrite: j.overwrite,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
ImagePreview: j.input.MarkerImagePreviews,
|
|
Screenshot: j.input.MarkerScreenshots,
|
|
|
|
generator: g,
|
|
}
|
|
|
|
markers := task.markersNeeded(ctx)
|
|
if markers > 0 {
|
|
j.totals.markers += int64(markers)
|
|
j.totals.tasks++
|
|
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
if j.input.Transcodes {
|
|
forceTranscode := j.input.ForceTranscodes
|
|
task := &GenerateTranscodeTask{
|
|
Scene: *scene,
|
|
Overwrite: j.overwrite,
|
|
Force: forceTranscode,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
g: g,
|
|
}
|
|
if task.required() {
|
|
j.totals.transcodes++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
if j.input.Phashes {
|
|
// generate for all files in scene
|
|
for _, f := range scene.Files.List() {
|
|
task := &GeneratePhashTask{
|
|
repository: r,
|
|
File: f,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
Overwrite: j.overwrite,
|
|
}
|
|
|
|
if task.required() {
|
|
j.totals.phashes++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
}
|
|
|
|
if j.input.InteractiveHeatmapsSpeeds {
|
|
task := &GenerateInteractiveHeatmapSpeedTask{
|
|
repository: r,
|
|
Scene: *scene,
|
|
Overwrite: j.overwrite,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
}
|
|
|
|
if task.required() {
|
|
j.totals.interactiveHeatmapSpeeds++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
}
|
|
|
|
func (j *GenerateJob) queueMarkerJob(g *generate.Generator, marker *models.SceneMarker, queue chan<- Task) {
|
|
task := &GenerateMarkersTask{
|
|
repository: j.repository,
|
|
Marker: marker,
|
|
Overwrite: j.overwrite,
|
|
fileNamingAlgorithm: j.fileNamingAlgo,
|
|
generator: g,
|
|
}
|
|
j.totals.markers++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
|
|
func (j *GenerateJob) queueImageJob(g *generate.Generator, image *models.Image, queue chan<- Task) {
|
|
if j.input.ImageThumbnails {
|
|
task := &GenerateImageThumbnailTask{
|
|
Image: *image,
|
|
Overwrite: j.overwrite,
|
|
}
|
|
|
|
if task.required() {
|
|
j.totals.imageThumbnails++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
|
|
if j.input.ClipPreviews {
|
|
task := &GenerateClipPreviewTask{
|
|
Image: *image,
|
|
Overwrite: j.overwrite,
|
|
}
|
|
|
|
if task.required() {
|
|
j.totals.clipPreviews++
|
|
j.totals.tasks++
|
|
queue <- task
|
|
}
|
|
}
|
|
}
|