353 lines
9.2 KiB
Go
353 lines
9.2 KiB
Go
package hailuo
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/pkg/errors"
|
|
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/relay/channel"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/service"
|
|
)
|
|
|
|
type TaskAdaptor struct {
|
|
ChannelType int
|
|
apiKey string
|
|
baseURL string
|
|
}
|
|
|
|
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
|
|
a.ChannelType = info.ChannelType
|
|
a.baseURL = info.ChannelBaseUrl
|
|
a.apiKey = info.ApiKey
|
|
}
|
|
|
|
func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
|
|
return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionGenerate)
|
|
}
|
|
|
|
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
|
return fmt.Sprintf("%s%s", a.baseURL, TextToVideoEndpoint), nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+a.apiKey)
|
|
return nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
|
|
v, exists := c.Get("task_request")
|
|
if !exists {
|
|
return nil, fmt.Errorf("request not found in context")
|
|
}
|
|
req, ok := v.(relaycommon.TaskSubmitReq)
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid request type in context")
|
|
}
|
|
|
|
body, err := a.convertToRequestPayload(&req)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "convert request payload failed")
|
|
}
|
|
|
|
data, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return bytes.NewReader(data), nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
|
|
return channel.DoTaskApiRequest(a, c, info, requestBody)
|
|
}
|
|
|
|
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
_ = resp.Body.Close()
|
|
|
|
var hResp TextToVideoResponse
|
|
if err := json.Unmarshal(responseBody, &hResp); err != nil {
|
|
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
if hResp.BaseResp.StatusCode != StatusSuccess {
|
|
taskErr = service.TaskErrorWrapper(
|
|
fmt.Errorf("hailuo api error: %s", hResp.BaseResp.StatusMsg),
|
|
strconv.Itoa(hResp.BaseResp.StatusCode),
|
|
http.StatusBadRequest,
|
|
)
|
|
return
|
|
}
|
|
|
|
ov := dto.NewOpenAIVideo()
|
|
ov.ID = hResp.TaskID
|
|
ov.TaskID = hResp.TaskID
|
|
ov.CreatedAt = time.Now().Unix()
|
|
ov.Model = info.OriginModelName
|
|
|
|
c.JSON(http.StatusOK, ov)
|
|
return hResp.TaskID, responseBody, nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any) (*http.Response, error) {
|
|
taskID, ok := body["task_id"].(string)
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid task_id")
|
|
}
|
|
|
|
uri := fmt.Sprintf("%s%s?task_id=%s", baseUrl, QueryTaskEndpoint, taskID)
|
|
|
|
req, err := http.NewRequest(http.MethodGet, uri, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Accept", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+key)
|
|
|
|
return service.GetHttpClient().Do(req)
|
|
}
|
|
|
|
func (a *TaskAdaptor) GetModelList() []string {
|
|
return ModelList
|
|
}
|
|
|
|
func (a *TaskAdaptor) GetChannelName() string {
|
|
return ChannelName
|
|
}
|
|
|
|
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*TextToVideoRequest, error) {
|
|
modelConfig := GetModelConfig(req.Model)
|
|
if !contains(ModelList, req.Model) {
|
|
return nil, fmt.Errorf("unsupported model: %s", req.Model)
|
|
}
|
|
|
|
duration := DefaultDuration
|
|
if req.Duration > 0 {
|
|
duration = req.Duration
|
|
}
|
|
|
|
if !containsInt(modelConfig.SupportedDurations, duration) {
|
|
return nil, fmt.Errorf("duration %d is not supported by model %s, supported durations: %v",
|
|
duration, req.Model, modelConfig.SupportedDurations)
|
|
}
|
|
|
|
resolution := modelConfig.DefaultResolution
|
|
if req.Size != "" {
|
|
resolution = a.parseResolutionFromSize(req.Size, modelConfig)
|
|
}
|
|
|
|
if !contains(modelConfig.SupportedResolutions, resolution) {
|
|
return nil, fmt.Errorf("resolution %s is not supported by model %s, supported resolutions: %v",
|
|
resolution, req.Model, modelConfig.SupportedResolutions)
|
|
}
|
|
|
|
hailuoReq := &TextToVideoRequest{
|
|
Model: req.Model,
|
|
Prompt: req.Prompt,
|
|
Duration: &duration,
|
|
Resolution: resolution,
|
|
}
|
|
|
|
promptOptimizer := DefaultPromptOptimizer
|
|
hailuoReq.PromptOptimizer = &promptOptimizer
|
|
|
|
metadata := req.Metadata
|
|
if metadata != nil {
|
|
metadataBytes, err := json.Marshal(metadata)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "marshal metadata failed")
|
|
}
|
|
|
|
var metadataMap map[string]interface{}
|
|
if err := json.Unmarshal(metadataBytes, &metadataMap); err != nil {
|
|
return nil, errors.Wrap(err, "unmarshal metadata failed")
|
|
}
|
|
|
|
if val, exists := metadataMap["prompt_optimizer"]; exists {
|
|
if boolVal, ok := val.(bool); ok {
|
|
hailuoReq.PromptOptimizer = &boolVal
|
|
}
|
|
}
|
|
|
|
if modelConfig.HasFastPretreatment {
|
|
if val, exists := metadataMap["fast_pretreatment"]; exists {
|
|
if boolVal, ok := val.(bool); ok {
|
|
hailuoReq.FastPretreatment = &boolVal
|
|
}
|
|
}
|
|
}
|
|
|
|
if val, exists := metadataMap["callback_url"]; exists {
|
|
if strVal, ok := val.(string); ok {
|
|
hailuoReq.CallbackURL = strVal
|
|
}
|
|
}
|
|
|
|
if val, exists := metadataMap["aigc_watermark"]; exists {
|
|
if boolVal, ok := val.(bool); ok {
|
|
hailuoReq.AigcWatermark = &boolVal
|
|
}
|
|
}
|
|
}
|
|
|
|
if req.HasImage() {
|
|
return nil, fmt.Errorf("image input is not supported by hailuo video generation")
|
|
}
|
|
|
|
return hailuoReq, nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConfig) string {
|
|
switch {
|
|
case strings.Contains(size, "1080"):
|
|
return Resolution1080P
|
|
case strings.Contains(size, "768"):
|
|
return Resolution768P
|
|
case strings.Contains(size, "720"):
|
|
return Resolution720P
|
|
default:
|
|
return modelConfig.DefaultResolution
|
|
}
|
|
}
|
|
|
|
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
|
|
resTask := QueryTaskResponse{}
|
|
if err := json.Unmarshal(respBody, &resTask); err != nil {
|
|
return nil, errors.Wrap(err, "unmarshal task result failed")
|
|
}
|
|
|
|
taskResult := relaycommon.TaskInfo{}
|
|
|
|
if resTask.BaseResp.StatusCode == StatusSuccess {
|
|
taskResult.Code = 0
|
|
} else {
|
|
taskResult.Code = resTask.BaseResp.StatusCode
|
|
taskResult.Reason = resTask.BaseResp.StatusMsg
|
|
taskResult.Status = model.TaskStatusFailure
|
|
taskResult.Progress = "100%"
|
|
}
|
|
|
|
switch resTask.Status {
|
|
case TaskStatusPreparing, TaskStatusQueueing, TaskStatusProcessing:
|
|
taskResult.Status = model.TaskStatusInProgress
|
|
taskResult.Progress = "30%"
|
|
if resTask.Status == TaskStatusProcessing {
|
|
taskResult.Progress = "50%"
|
|
}
|
|
case TaskStatusSuccess:
|
|
taskResult.Status = model.TaskStatusSuccess
|
|
taskResult.Progress = "100%"
|
|
taskResult.Url = a.buildVideoURL(resTask.TaskID, resTask.FileID)
|
|
case TaskStatusFailed:
|
|
taskResult.Status = model.TaskStatusFailure
|
|
taskResult.Progress = "100%"
|
|
if taskResult.Reason == "" {
|
|
taskResult.Reason = "task failed"
|
|
}
|
|
default:
|
|
taskResult.Status = model.TaskStatusInProgress
|
|
taskResult.Progress = "30%"
|
|
}
|
|
|
|
return &taskResult, nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
|
|
var hailuoResp QueryTaskResponse
|
|
if err := json.Unmarshal(originTask.Data, &hailuoResp); err != nil {
|
|
return nil, errors.Wrap(err, "unmarshal hailuo task data failed")
|
|
}
|
|
|
|
openAIVideo := originTask.ToOpenAIVideo()
|
|
if hailuoResp.BaseResp.StatusCode != StatusSuccess {
|
|
openAIVideo.Error = &dto.OpenAIVideoError{
|
|
Message: hailuoResp.BaseResp.StatusMsg,
|
|
Code: strconv.Itoa(hailuoResp.BaseResp.StatusCode),
|
|
}
|
|
}
|
|
|
|
jsonData, err := common.Marshal(openAIVideo)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "marshal openai video failed")
|
|
}
|
|
|
|
return jsonData, nil
|
|
}
|
|
|
|
func (a *TaskAdaptor) buildVideoURL(_, fileID string) string {
|
|
if a.apiKey == "" || a.baseURL == "" {
|
|
return ""
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/v1/files/retrieve?file_id=%s", a.baseURL, fileID)
|
|
|
|
req, err := http.NewRequest(http.MethodGet, url, nil)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
|
|
req.Header.Set("Accept", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+a.apiKey)
|
|
|
|
resp, err := service.GetHttpClient().Do(req)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
|
|
var retrieveResp RetrieveFileResponse
|
|
if err := json.Unmarshal(responseBody, &retrieveResp); err != nil {
|
|
return ""
|
|
}
|
|
|
|
if retrieveResp.BaseResp.StatusCode != StatusSuccess {
|
|
return ""
|
|
}
|
|
|
|
return retrieveResp.File.DownloadURL
|
|
}
|
|
|
|
func contains(slice []string, item string) bool {
|
|
for _, s := range slice {
|
|
if s == item {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func containsInt(slice []int, item int) bool {
|
|
for _, s := range slice {
|
|
if s == item {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|