LocalAI/core/http/endpoints/localai/video.go

package localai

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/labstack/echo/v4"

	"github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/core/http/middleware"
	"github.com/mudler/LocalAI/core/schema"

	"github.com/mudler/LocalAI/core/backend"

	"github.com/mudler/xlog"

	"github.com/mudler/LocalAI/pkg/httpclient"
	model "github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/utils"
)

// Downloading user-supplied media URLs legitimately follows redirects (CDNs);
// WithFollowRedirects still strips any credential header on a cross-host hop.
var videoDownloadClient = httpclient.NewWithTimeout(30*time.Second, httpclient.WithFollowRedirects())

func downloadFile(url string) (string, error) {
	if err := utils.ValidateExternalURL(url); err != nil {
		return "", fmt.Errorf("URL validation failed: %w", err)
	}

	// Get the data
	resp, err := videoDownloadClient.Get(url)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Create the file
	out, err := os.CreateTemp("", "video")
	if err != nil {
		return "", err
	}
	defer out.Close()

	// Write the body to file
	_, err = io.Copy(out, resp.Body)
	return out.Name(), err
}

//

/*
*

	curl http://localhost:8080/v1/images/generations \
	  -H "Content-Type: application/json" \
	  -d '{
	    "prompt": "A cute baby sea otter",
	    "n": 1,
	    "size": "512x512"
	  }'

*
*/
// VideoEndpoint
// @Summary Creates a video given a prompt.
// @Tags video
// @Param request body schema.VideoRequest true "query params"
// @Success 200 {object} schema.OpenAIResponse "Response"
// @Router /video [post]
func VideoEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc {
	return func(c echo.Context) error {
		input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.VideoRequest)
		if !ok || input.Model == "" {
			xlog.Error("Video Endpoint - Invalid Input")
			return echo.ErrBadRequest
		}

		config, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig)
		if !ok || config == nil {
			xlog.Error("Video Endpoint - Invalid Config")
			return echo.ErrBadRequest
		}

		// Stage a base64- or URL-provided image into a temp file so the
		// backend can read it as a path. Used for both start_image and
		// (optional) end_image. Returns the temp file path, or "" if the
		// input is empty. Caller is responsible for the defer-cleanup.
		stageImage := func(ref string) (string, error) {
			if ref == "" {
				return "", nil
			}
			var fileData []byte
			var err error
			if strings.HasPrefix(ref, "http://") || strings.HasPrefix(ref, "https://") {
				out, derr := downloadFile(ref)
				if derr != nil {
					return "", fmt.Errorf("failed downloading file: %w", derr)
				}
				defer os.RemoveAll(out)
				fileData, err = os.ReadFile(out)
				if err != nil {
					return "", fmt.Errorf("failed reading file: %w", err)
				}
			} else {
				fileData, err = base64.StdEncoding.DecodeString(ref)
				if err != nil {
					return "", err
				}
			}
			outputFile, err := os.CreateTemp(appConfig.GeneratedContentDir, "b64")
			if err != nil {
				return "", err
			}
			writer := bufio.NewWriter(outputFile)
			if _, err := writer.Write(fileData); err != nil {
				outputFile.Close()
				return "", err
			}
			if err := writer.Flush(); err != nil {
				outputFile.Close()
				return "", err
			}
			outputFile.Close()
			return outputFile.Name(), nil
		}

		src, err := stageImage(input.StartImage)
		if err != nil {
			return err
		}
		if src != "" {
			defer os.RemoveAll(src)
		}

		endSrc, err := stageImage(input.EndImage)
		if err != nil {
			return err
		}
		if endSrc != "" {
			defer os.RemoveAll(endSrc)
		}

		xlog.Debug("Parameter Config", "config", config)

		switch config.Backend {
		case "stablediffusion":
			config.Backend = model.StableDiffusionGGMLBackend
		case "":
			config.Backend = model.StableDiffusionGGMLBackend
		}

		width := input.Width
		height := input.Height

		if width == 0 {
			width = 512
		}
		if height == 0 {
			height = 512
		}

		b64JSON := input.ResponseFormat == "b64_json"

		tempDir := ""
		if !b64JSON {
			tempDir = filepath.Join(appConfig.GeneratedContentDir, "videos")
		}
		// Create a temporary file
		outputFile, err := os.CreateTemp(tempDir, "b64")
		if err != nil {
			return err
		}
		outputFile.Close()

		// TODO: use mime type to determine the extension
		output := outputFile.Name() + ".mp4"

		// Rename the temporary file
		err = os.Rename(outputFile.Name(), output)
		if err != nil {
			return err
		}

		baseURL := middleware.BaseURL(c)

		xlog.Debug("VideoEndpoint: Calling VideoGeneration",
			"num_frames", input.NumFrames,
			"fps", input.FPS,
			"cfg_scale", input.CFGScale,
			"step", input.Step,
			"seed", input.Seed,
			"width", width,
			"height", height,
			"negative_prompt", input.NegativePrompt)

		fn, err := backend.VideoGeneration(
			height,
			width,
			input.Prompt,
			input.NegativePrompt,
			src,
			endSrc,
			output,
			input.NumFrames,
			input.FPS,
			input.Seed,
			input.CFGScale,
			input.Step,
			ml,
			*config,
			appConfig,
		)
		if err != nil {
			return err
		}
		if err := fn(); err != nil {
			return err
		}

		item := &schema.Item{}

		if b64JSON {
			defer os.RemoveAll(output)
			data, err := os.ReadFile(output)
			if err != nil {
				return err
			}
			item.B64JSON = base64.StdEncoding.EncodeToString(data)
		} else {
			base := filepath.Base(output)
			item.URL, err = url.JoinPath(baseURL, "generated-videos", base)
			if err != nil {
				return err
			}
		}

		id := uuid.New().String()
		created := int(time.Now().Unix())
		resp := &schema.OpenAIResponse{
			ID:      id,
			Created: created,
			Data:    []schema.Item{*item},
		}

		jsonResult, _ := json.Marshal(resp)
		xlog.Debug("Response", "response", string(jsonResult))

		// Return the prediction in the response body
		return c.JSON(200, resp)
	}
}