mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-12-26 18:15:57 +08:00
merge v4.2.1
This commit is contained in:
@@ -34,10 +34,8 @@ type Service struct {
|
||||
db *gorm.DB
|
||||
uploadManager *oss.UploaderManager
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
userService *service.UserService
|
||||
wsService *service.WebsocketService
|
||||
clientIds map[uint]string
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService, wsService *service.WebsocketService) *Service {
|
||||
@@ -45,11 +43,9 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
|
||||
httpClient: req.C().SetTimeout(time.Minute * 3),
|
||||
db: db,
|
||||
taskQueue: store.NewRedisQueue("DallE_Task_Queue", redisCli),
|
||||
notifyQueue: store.NewRedisQueue("DallE_Notify_Queue", redisCli),
|
||||
wsService: wsService,
|
||||
uploadManager: manager,
|
||||
userService: userService,
|
||||
clientIds: map[uint]string{},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -60,7 +56,7 @@ func (s *Service) PushTask(task types.DallTask) {
|
||||
}
|
||||
|
||||
func (s *Service) Run() {
|
||||
// 将数据库中未提交的人物加载到队列
|
||||
// 将数据库中未提交的任务加载到队列
|
||||
var jobs []model.DallJob
|
||||
s.db.Where("progress", 0).Find(&jobs)
|
||||
for _, v := range jobs {
|
||||
@@ -84,16 +80,16 @@ func (s *Service) Run() {
|
||||
continue
|
||||
}
|
||||
logger.Infof("handle a new DALL-E task: %+v", task)
|
||||
s.clientIds[task.Id] = task.ClientId
|
||||
_, err = s.Image(task, false)
|
||||
if err != nil {
|
||||
logger.Errorf("error with image task: %v", err)
|
||||
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress,
|
||||
"err_msg": err.Error(),
|
||||
})
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
|
||||
}
|
||||
go func() {
|
||||
_, err = s.Image(task, false)
|
||||
if err != nil {
|
||||
logger.Errorf("error with image task: %v", err)
|
||||
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress,
|
||||
"err_msg": err.Error(),
|
||||
})
|
||||
}
|
||||
}()
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -212,10 +208,9 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
|
||||
return "", fmt.Errorf("err with update database: %v", err)
|
||||
}
|
||||
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
|
||||
var content string
|
||||
if sync {
|
||||
imgURL, err := s.downloadImage(task.Id, int(task.UserId), res.Data[0].Url)
|
||||
imgURL, err := s.downloadImage(task.Id, res.Data[0].Url)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error with download image: %v", err)
|
||||
}
|
||||
@@ -225,26 +220,6 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
|
||||
return content, nil
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
logger.Info("Running DALL-E task notify checking ...")
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Debugf("notify message: %+v", message)
|
||||
client := s.wsService.Clients.Get(message.ClientId)
|
||||
if client == nil {
|
||||
continue
|
||||
}
|
||||
utils.SendChannelMsg(client, types.ChDall, message.Message)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskStatus() {
|
||||
go func() {
|
||||
logger.Info("Running DALL-E task status checking ...")
|
||||
@@ -254,7 +229,7 @@ func (s *Service) CheckTaskStatus() {
|
||||
s.db.Where("progress < ?", 100).Find(&jobs)
|
||||
for _, job := range jobs {
|
||||
// 超时的任务标记为失败
|
||||
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
|
||||
if time.Since(job.CreatedAt) > time.Minute*10 {
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = "任务超时"
|
||||
s.db.Updates(&job)
|
||||
@@ -301,7 +276,7 @@ func (s *Service) DownloadImages() {
|
||||
}
|
||||
|
||||
logger.Infof("try to download image: %s", v.OrgURL)
|
||||
imgURL, err := s.downloadImage(v.Id, int(v.UserId), v.OrgURL)
|
||||
imgURL, err := s.downloadImage(v.Id, v.OrgURL)
|
||||
if err != nil {
|
||||
logger.Error("error with download image: %s, error: %v", imgURL, err)
|
||||
continue
|
||||
@@ -316,7 +291,7 @@ func (s *Service) DownloadImages() {
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string, error) {
|
||||
func (s *Service) downloadImage(jobId uint, orgURL string) (string, error) {
|
||||
// sava image
|
||||
imgURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(orgURL, false)
|
||||
if err != nil {
|
||||
@@ -328,6 +303,5 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
|
||||
if res.Error != nil {
|
||||
return "", err
|
||||
}
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[jobId], UserId: userId, JobId: int(jobId), Message: service.TaskStatusFinished})
|
||||
return imgURL, nil
|
||||
}
|
||||
|
||||
52
api/service/migration_service.go
Normal file
52
api/service/migration_service.go
Normal file
@@ -0,0 +1,52 @@
|
||||
package service
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"geekai/store/model"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type MigrationService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewMigrationService(db *gorm.DB) *MigrationService {
|
||||
return &MigrationService{db: db}
|
||||
}
|
||||
|
||||
func (s *MigrationService) Migrate() error {
|
||||
err := s.db.AutoMigrate(
|
||||
&model.AdminUser{},
|
||||
&model.ApiKey{},
|
||||
&model.AppType{},
|
||||
&model.ChatItem{},
|
||||
&model.ChatMessage{},
|
||||
&model.ChatModel{},
|
||||
&model.ChatRole{},
|
||||
&model.Config{},
|
||||
&model.DallJob{},
|
||||
&model.File{},
|
||||
&model.Function{},
|
||||
&model.InviteCode{},
|
||||
&model.InviteLog{},
|
||||
&model.Menu{},
|
||||
&model.MidJourneyJob{},
|
||||
&model.Order{},
|
||||
&model.PowerLog{},
|
||||
&model.Product{},
|
||||
&model.Redeem{},
|
||||
&model.SdJob{},
|
||||
&model.SunoJob{},
|
||||
&model.User{},
|
||||
&model.UserLoginLog{},
|
||||
&model.VideoJob{},
|
||||
)
|
||||
return err
|
||||
}
|
||||
@@ -15,10 +15,11 @@ import (
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-redis/redis/v8"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -26,23 +27,19 @@ import (
|
||||
type Service struct {
|
||||
client *Client // MJ Client
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
db *gorm.DB
|
||||
wsService *service.WebsocketService
|
||||
uploaderManager *oss.UploaderManager
|
||||
userService *service.UserService
|
||||
clientIds map[uint]string
|
||||
}
|
||||
|
||||
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager, wsService *service.WebsocketService, userService *service.UserService) *Service {
|
||||
return &Service{
|
||||
db: db,
|
||||
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
|
||||
notifyQueue: store.NewRedisQueue("MidJourney_Notify_Queue", redisCli),
|
||||
client: client,
|
||||
wsService: wsService,
|
||||
uploaderManager: manager,
|
||||
clientIds: map[uint]string{},
|
||||
userService: userService,
|
||||
}
|
||||
}
|
||||
@@ -59,7 +56,6 @@ func (s *Service) Run() {
|
||||
continue
|
||||
}
|
||||
task.Id = v.Id
|
||||
s.clientIds[task.Id] = task.ClientId
|
||||
s.PushTask(task)
|
||||
}
|
||||
|
||||
@@ -96,7 +92,6 @@ func (s *Service) Run() {
|
||||
if task.Mode == "" {
|
||||
task.Mode = "fast"
|
||||
}
|
||||
s.clientIds[task.Id] = task.ClientId
|
||||
|
||||
var job model.MidJourneyJob
|
||||
tx := s.db.Where("id = ?", task.Id).First(&job)
|
||||
@@ -139,7 +134,6 @@ func (s *Service) Run() {
|
||||
// update the task progress
|
||||
s.db.Updates(&job)
|
||||
// 任务失败,通知前端
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
logger.Infof("任务提交成功:%+v", res)
|
||||
@@ -178,24 +172,6 @@ func GetImageHash(action string) string {
|
||||
return split[len(split)-1]
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
logger.Debugf("receive a new mj notify message: %+v", message)
|
||||
client := s.wsService.Clients.Get(message.ClientId)
|
||||
if client == nil {
|
||||
continue
|
||||
}
|
||||
utils.SendChannelMsg(client, types.ChMj, message.Message)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) DownloadImages() {
|
||||
go func() {
|
||||
var items []model.MidJourneyJob
|
||||
@@ -228,12 +204,6 @@ func (s *Service) DownloadImages() {
|
||||
|
||||
v.ImgURL = imgURL
|
||||
s.db.Updates(&v)
|
||||
|
||||
s.notifyQueue.RPush(service.NotifyMessage{
|
||||
ClientId: s.clientIds[v.Id],
|
||||
UserId: v.UserId,
|
||||
JobId: int(v.Id),
|
||||
Message: service.TaskStatusFinished})
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 5)
|
||||
@@ -259,7 +229,7 @@ func (s *Service) SyncTaskProgress() {
|
||||
|
||||
for _, job := range jobs {
|
||||
// 10 分钟还没完成的任务标记为失败
|
||||
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
|
||||
if time.Since(job.CreatedAt) > time.Minute*10 {
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = "任务超时"
|
||||
s.db.Updates(&job)
|
||||
@@ -283,18 +253,12 @@ func (s *Service) SyncTaskProgress() {
|
||||
"err_msg": task.FailReason,
|
||||
})
|
||||
logger.Errorf("task failed: %v", task.FailReason)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{
|
||||
ClientId: s.clientIds[job.Id],
|
||||
UserId: job.UserId,
|
||||
JobId: int(job.Id),
|
||||
Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
|
||||
if len(task.Buttons) > 0 {
|
||||
job.Hash = GetImageHash(task.Buttons[0].CustomId)
|
||||
}
|
||||
oldProgress := job.Progress
|
||||
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
|
||||
if task.ImageUrl != "" {
|
||||
job.OrgURL = task.ImageUrl
|
||||
@@ -304,19 +268,6 @@ func (s *Service) SyncTaskProgress() {
|
||||
logger.Errorf("error with update database: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 通知前端更新任务进度
|
||||
if oldProgress != job.Progress {
|
||||
message := service.TaskStatusRunning
|
||||
if job.Progress == 100 {
|
||||
message = service.TaskStatusFinished
|
||||
}
|
||||
s.notifyQueue.RPush(service.NotifyMessage{
|
||||
ClientId: s.clientIds[job.Id],
|
||||
UserId: job.UserId,
|
||||
JobId: int(job.Id),
|
||||
Message: message})
|
||||
}
|
||||
}
|
||||
|
||||
// 找出失败的任务,并恢复其扣减算力
|
||||
|
||||
@@ -16,9 +16,10 @@ import (
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"time"
|
||||
|
||||
"github.com/go-redis/redis/v8"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -30,7 +31,6 @@ var logger = logger2.GetLogger()
|
||||
type Service struct {
|
||||
httpClient *req.Client
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
db *gorm.DB
|
||||
uploadManager *oss.UploaderManager
|
||||
wsService *service.WebsocketService
|
||||
@@ -41,7 +41,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelD
|
||||
return &Service{
|
||||
httpClient: req.C(),
|
||||
taskQueue: store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli),
|
||||
notifyQueue: store.NewRedisQueue("StableDiffusion_Queue", redisCli),
|
||||
db: db,
|
||||
wsService: wsService,
|
||||
uploadManager: manager,
|
||||
@@ -102,8 +101,6 @@ func (s *Service) Run() {
|
||||
"progress": service.FailTaskProgress,
|
||||
"err_msg": err.Error(),
|
||||
})
|
||||
// 通知前端,任务失败
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -225,15 +222,12 @@ func (s *Service) Txt2Img(task types.SdTask) error {
|
||||
|
||||
// task finished
|
||||
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFinished})
|
||||
return nil
|
||||
default:
|
||||
err, resp := s.checkTaskProgress(apiKey)
|
||||
resp, err := s.checkTaskProgress(apiKey)
|
||||
// 更新任务进度
|
||||
if err == nil && resp.Progress > 0 {
|
||||
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
|
||||
// 发送更新状态信号
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusRunning})
|
||||
}
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
@@ -242,7 +236,7 @@ func (s *Service) Txt2Img(task types.SdTask) error {
|
||||
}
|
||||
|
||||
// 执行任务
|
||||
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressResp) {
|
||||
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (*TaskProgressResp, error) {
|
||||
apiURL := fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", apiKey.ApiURL)
|
||||
var res TaskProgressResp
|
||||
response, err := s.httpClient.R().
|
||||
@@ -250,13 +244,13 @@ func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressRe
|
||||
SetSuccessResult(&res).
|
||||
Get(apiURL)
|
||||
if err != nil {
|
||||
return err, nil
|
||||
return nil, err
|
||||
}
|
||||
if response.IsErrorState() {
|
||||
return fmt.Errorf("error http code status: %v", response.Status), nil
|
||||
return nil, fmt.Errorf("error http code status: %v", response.Status)
|
||||
}
|
||||
|
||||
return nil, &res
|
||||
return &res, nil
|
||||
}
|
||||
|
||||
func (s *Service) PushTask(task types.SdTask) {
|
||||
@@ -264,25 +258,6 @@ func (s *Service) PushTask(task types.SdTask) {
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
logger.Info("Running Stable-Diffusion task notify checking ...")
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
logger.Debugf("notify message: %+v", message)
|
||||
client := s.wsService.Clients.Get(message.ClientId)
|
||||
if client == nil {
|
||||
continue
|
||||
}
|
||||
utils.SendChannelMsg(client, types.ChSd, message.Message)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// CheckTaskStatus 检查任务状态,自动删除过期或者失败的任务
|
||||
func (s *Service) CheckTaskStatus() {
|
||||
go func() {
|
||||
@@ -297,7 +272,7 @@ func (s *Service) CheckTaskStatus() {
|
||||
|
||||
for _, job := range jobs {
|
||||
// 5 分钟还没完成的任务标记为失败
|
||||
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
|
||||
if time.Since(job.CreatedAt) > time.Minute*5 {
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = "任务超时"
|
||||
s.db.Updates(&job)
|
||||
|
||||
@@ -18,10 +18,11 @@ import (
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/go-redis/redis/v8"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -35,7 +36,6 @@ type Service struct {
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
wsService *service.WebsocketService
|
||||
clientIds map[string]string
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
@@ -47,7 +47,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
|
||||
notifyQueue: store.NewRedisQueue("Suno_Notify_Queue", redisCli),
|
||||
uploadManager: manager,
|
||||
wsService: wsService,
|
||||
clientIds: map[string]string{},
|
||||
userService: userService,
|
||||
}
|
||||
}
|
||||
@@ -70,7 +69,6 @@ func (s *Service) Run() {
|
||||
}
|
||||
task.Id = v.Id
|
||||
s.PushTask(task)
|
||||
s.clientIds[v.TaskId] = task.ClientId
|
||||
}
|
||||
logger.Info("Starting Suno job consumer...")
|
||||
go func() {
|
||||
@@ -95,7 +93,6 @@ func (s *Service) Run() {
|
||||
"err_msg": err.Error(),
|
||||
"progress": service.FailTaskProgress,
|
||||
})
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -104,7 +101,6 @@ func (s *Service) Run() {
|
||||
"task_id": r.Data,
|
||||
"channel": r.Channel,
|
||||
})
|
||||
s.clientIds[r.Data] = task.ClientId
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -138,7 +134,7 @@ func (s *Service) Create(task types.SunoTask) (RespVo, error) {
|
||||
if task.Type == 1 {
|
||||
reqBody["gpt_description_prompt"] = task.Prompt
|
||||
} else { // 自定义模式
|
||||
reqBody["prompt"] = task.Prompt
|
||||
reqBody["prompt"] = task.Lyrics
|
||||
reqBody["tags"] = task.Tags
|
||||
reqBody["mv"] = task.Model
|
||||
reqBody["title"] = task.Title
|
||||
@@ -146,7 +142,7 @@ func (s *Service) Create(task types.SunoTask) (RespVo, error) {
|
||||
|
||||
var res RespVo
|
||||
apiURL := fmt.Sprintf("%s/suno/submit/music", apiKey.ApiURL)
|
||||
logger.Debugf("API URL: %s, request body: %+v", apiURL, reqBody)
|
||||
logger.Debugf("API URL: %s, request body: %s", apiURL, utils.JsonEncode(reqBody))
|
||||
r, err := req.C().R().
|
||||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetBody(reqBody).
|
||||
@@ -262,27 +258,6 @@ func (s *Service) Upload(task types.SunoTask) (RespVo, error) {
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
logger.Info("Running Suno task notify checking ...")
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
logger.Debugf("notify message: %+v", message)
|
||||
logger.Debugf("client id: %+v", s.wsService.Clients)
|
||||
client := s.wsService.Clients.Get(message.ClientId)
|
||||
logger.Debugf("%+v", client)
|
||||
if client == nil {
|
||||
continue
|
||||
}
|
||||
utils.SendChannelMsg(client, types.ChSuno, message.Message)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) DownloadFiles() {
|
||||
go func() {
|
||||
var items []model.SunoJob
|
||||
@@ -311,7 +286,6 @@ func (s *Service) DownloadFiles() {
|
||||
v.AudioURL = audioURL
|
||||
v.Progress = 100
|
||||
s.db.Updates(&v)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.TaskId], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished})
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 10)
|
||||
@@ -377,12 +351,10 @@ func (s *Service) SyncTaskProgress() {
|
||||
}
|
||||
}
|
||||
tx.Commit()
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFinished})
|
||||
} else if task.Data.FailReason != "" {
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = task.Data.FailReason
|
||||
s.db.Updates(&job)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -12,109 +12,89 @@ type NotifyMessage struct {
|
||||
ClientId string `json:"client_id"`
|
||||
JobId int `json:"job_id"`
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
const TranslatePromptTemplate = "Translate the following painting prompt words into English keyword phrases. Without any explanation, directly output the keyword phrases separated by commas. The content to be translated is: [%s]"
|
||||
|
||||
const ImagePromptOptimizeTemplate = `
|
||||
Create a highly effective prompt to provide to an AI image generation tool in order to create an artwork based on a desired concept.
|
||||
以下是一条 AI 提示词示例,用于优化和扩写绘图提示词:
|
||||
|
||||
Please specify details about the artwork, such as the style, subject, mood, and other important characteristics you want the resulting image to have.
|
||||
请你作为一名专业的 AI 绘图提示词优化专家,基于用户提供的简单绘图描述,生成一份详细、专业且富有创意的 AI 绘图提示词指令。在优化过程中,你需要做到以下几点:
|
||||
|
||||
Remember, prompts should always be output in English.
|
||||
1. 深入理解用户描述的核心意图和关键元素,挖掘潜在的细节和情感氛围,将其融入到提示词中。
|
||||
2. 丰富画面细节,包括但不限于场景背景、人物特征、物体属性、光影效果、色彩搭配等,使画面更加生动逼真。
|
||||
3. 运用专业的艺术风格术语,如超现实主义、印象派、赛博朋克等,为画面增添独特的艺术魅力。
|
||||
4. 考虑构图和视角,如俯视、仰视、特写、全景等,提升画面的视觉冲击力。
|
||||
5. 确保提示词指令清晰、准确、完整,便于 AI 绘图模型理解和生成高质量图像。最终输出的提示词应简洁明了,避免冗余信息,以逗号分隔各个元素,突出重点,
|
||||
让用户能够直接复制使用,从而帮助用户将简单的想法转化为精美绝伦的画作。
|
||||
6. 不管用户输入的是什么语言,你务必要用英文输出优化后的提示词。
|
||||
7. 直接输出优化后的提示词,不要输出其他任何五官内容。
|
||||
|
||||
# Steps
|
||||
下面是一个提示词优化示例:
|
||||
===示例开始===
|
||||
原始指令 :一个穿着红色连衣裙的少女在花园里浇花,阳光明媚。
|
||||
|
||||
1. **Subject Description**: Describe the main subject of the image clearly. Include as much detail as possible about what should be in the scene. For example, "a majestic lion roaring at sunrise" or "a futuristic city with flying cars."
|
||||
|
||||
2. **Art Style**: Specify the art style you envision. Possible options include 'realistic', 'impressionist', a specific artist name, or imaginative styles like "cyberpunk." This helps the AI achieve your visual expectations.
|
||||
优化后的 AI 绘图提示词指令:一位年轻美丽的少女,约 16 - 18 岁,有着柔顺的黑色长发,披散在肩上,面容精致,眼神温柔而专注。她穿着一条复古风格的红色连衣裙,裙子上有精致的褶皱和白色的蕾丝花边,裙摆轻轻飘动。少女站在一个充满生机的花园中,花园里种满了各种各样的鲜花,有娇艳的玫瑰、淡雅的百合、缤纷的郁金香等,花朵色彩鲜艳,绿叶繁茂。她手持一个银色的 watering can(浇水壶),正在细心地给一朵盛开的玫瑰浇水。阳光从画面的右侧洒下,形成明亮而温暖的光晕,照亮了少女和整个花园,营造出一种宁静、美好的氛围,画面采用写实风格,光影效果逼真,色彩鲜明且富有层次感,构图以少女为中心,前景是盛开的花朵,背景是花园的树木和篱笆,整体画面充满诗意和浪漫气息。
|
||||
===示例结束===
|
||||
|
||||
3. **Mood or Atmosphere**: Convey the feeling you want the image to evoke. For instance, peaceful, chaotic, epic, etc.
|
||||
|
||||
4. **Color Palette and Lighting**: Mention color preferences or lighting. For example, "vibrant with shades of blue and purple" or "dim and dramatic lighting."
|
||||
|
||||
5. **Optional Features**: You can add any additional attributes, such as background details, attention to textures, or any specific kind of framing.
|
||||
|
||||
# Output Format
|
||||
|
||||
- **Prompt Format**: A descriptive phrase that includes key aspects of the artwork (subject, style, mood, colors, lighting, any optional features).
|
||||
|
||||
Here is an example of how the final prompt should look:
|
||||
|
||||
"An ethereal landscape featuring towering ice mountains, in an impressionist style reminiscent of Claude Monet, with a serene mood. The sky is glistening with soft purples and whites, with a gentle morning sun illuminating the scene."
|
||||
|
||||
**Please input the prompt words directly in English, and do not input any other explanatory statements**
|
||||
|
||||
# Examples
|
||||
|
||||
1. **Input**:
|
||||
- Subject: A white tiger in a dense jungle
|
||||
- Art Style: Realistic
|
||||
- Mood: Intense, mysterious
|
||||
- Lighting: Dramatic contrast with light filtering through leaves
|
||||
|
||||
**Output Prompt**: "A realistic rendering of a white tiger stealthily moving through a dense jungle, with an intense, mysterious mood. The lighting creates strong contrasts as beams of sunlight filter through a thick canopy of leaves."
|
||||
|
||||
2. **Input**:
|
||||
- Subject: An enchanted castle on a floating island
|
||||
- Art Style: Fantasy
|
||||
- Mood: Majestic, magical
|
||||
- Colors: Bright blues, greens, and gold
|
||||
|
||||
**Output Prompt**: "A majestic fantasy castle on a floating island above the clouds, with bright blues, greens, and golds to create a magical, dreamy atmosphere. Textured cobblestone details and glistening waters surround the scene."
|
||||
|
||||
# Notes
|
||||
|
||||
- Ensure that you mix different aspects to get a comprehensive and visually compelling prompt.
|
||||
- Be as descriptive as possible as it often helps generate richer, more detailed images.
|
||||
- If you want the image to resemble a particular artist's work, be sure to mention the artist explicitly. e.g., "in the style of Van Gogh."
|
||||
|
||||
The theme of the creation is:【%s】
|
||||
现在用户输入的原始提示词为:【%s】
|
||||
`
|
||||
|
||||
const LyricPromptTemplate = `
|
||||
你是一位才华横溢的作曲家,拥有丰富的情感和细腻的笔触,你对文字有着独特的感悟力,能将各种情感和意境巧妙地融入歌词中。
|
||||
请以【%s】为主题创作一首歌曲,歌曲时间不要太短,3分钟左右,不要输出任何解释性的内容。
|
||||
输出格式如下:
|
||||
下面是一个标准的歌词输出模板:
|
||||
歌曲名称
|
||||
第一节:
|
||||
{{歌词内容}}
|
||||
副歌:
|
||||
{{歌词内容}}
|
||||
|
||||
第二节:
|
||||
{{歌词内容}}
|
||||
副歌:
|
||||
{{歌词内容}}
|
||||
[Verse]
|
||||
[歌词]
|
||||
|
||||
尾声:
|
||||
{{歌词内容}}
|
||||
[Verse 2]
|
||||
[歌词]
|
||||
|
||||
[Chorus]
|
||||
[歌词]
|
||||
|
||||
[Verse 3]
|
||||
[歌词]
|
||||
|
||||
[Bridge]
|
||||
[歌词]
|
||||
|
||||
[Chorus]
|
||||
[歌词]
|
||||
|
||||
[Verse 4]
|
||||
[歌词]
|
||||
|
||||
[Bridge]
|
||||
假如此刻眼泪能倒流
|
||||
让我学会微笑不掩忧
|
||||
一次次的碎片堆积的愁
|
||||
最终也会开成希望的秋
|
||||
|
||||
[Chorus]
|
||||
假如我还能牵你的手
|
||||
天空也许会更蔚蓝悠游
|
||||
曾经那些未完成的错过
|
||||
愿能变成今天的收获
|
||||
`
|
||||
|
||||
const VideoPromptTemplate = `
|
||||
As an expert in video generation prompts, please create a detailed descriptive prompt for the following video concept. The description should include the setting, character appearance, actions, overall atmosphere, and camera angles. Please make it as detailed and vivid as possible to help ensure that every aspect of the video is accurately captured.
|
||||
const VideoPromptTemplate = `## 任务描述
|
||||
你是一位优秀AI视频创作专家,擅长编写专业的AI视频提示词,现在你的任务是对用户输入的简单视频描述提示词进行专业优化和扩写,使其转化为详细的、具备专业影视画面感的 AI 生成视频提示词指令。需涵盖风格、主体元素、环境氛围、细节特征、人物状态(若有)、镜头运用及整体氛围营造等方面,以生动形象、富有感染力且精准的描述,引导 AI 生成高质量的视频内容。下面是一个示例:
|
||||
===示例开始===
|
||||
输入: “汽车在沙漠功能上行驶”,
|
||||
输出: “纪实摄影风格,一辆尘土飞扬的复古越野车在无垠的沙漠公路上疾驰,车身线条硬朗,漆面斑驳,透露出岁月的痕迹。驾驶室内的司机戴着墨镜,专注地握着方向盘,眼神坚定地望向前方。夕阳的余晖洒在车身上,沙漠的沙丘在远处延绵起伏,一片金黄。广角镜头捕捉到车辆行驶时扬起的沙尘,营造出动感与冒险的氛围。远景全貌,强调速度感与环境辽阔。”
|
||||
===示例结束===
|
||||
|
||||
Please remember that regardless of the user’s input, the final output must be in English.
|
||||
## 输出要求:
|
||||
1. 直接输出扩写后的提示词就好,不要输出其他任何不相关信息
|
||||
2. 如果用户用中文提问,你就用中文回答,如果用英文提问,你也必须用英文回答。
|
||||
3. 请确保提示词的长度长度在1000个字以内。
|
||||
|
||||
# Details to Include
|
||||
|
||||
- Describe the overall visual style of the video (e.g., animated, realistic, retro tone, etc.)
|
||||
- Identify key characters or objects in the video and describe their appearance, attire, and expressions
|
||||
- Describe the environment of the scene, including weather, lighting, colors, and important details
|
||||
- Explain the behavior and interactions of the characters
|
||||
- Include any unique camera angles, movements, or special effects
|
||||
|
||||
# Output Format
|
||||
Provide the prompt in paragraph form, ensuring that the description is detailed enough for a video generation system to recreate the envisioned scene. Include the beginning, middle, and end of the scene to convey a complete storyline.
|
||||
|
||||
# Example
|
||||
**User Input:**
|
||||
“A small cat basking in the sun on a balcony.”
|
||||
|
||||
**Generated Prompt:**
|
||||
On a bright spring afternoon, an orange-striped kitten lies lazily on a balcony, basking in the warm sunlight. The iron railings around the balcony cast soft shadows that dance gently with the light. The cat’s eyes are half-closed, exuding a sense of contentment and tranquility in its surroundings. In the distance, a few fluffy white clouds drift slowly across the blue sky. The camera initially focuses on the cat’s face, capturing the delicate details of its fur, and then gradually zooms out to reveal the full balcony scene, immersing viewers in a moment of calm and relaxation.
|
||||
|
||||
The theme of the creation is:【%s】
|
||||
=====
|
||||
用户的输入的视频主题是:【%s】
|
||||
`
|
||||
|
||||
const MetaPromptTemplate = `
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
package video
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
logger2 "geekai/logger"
|
||||
"geekai/service"
|
||||
"geekai/service/oss"
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var logger = logger2.GetLogger()
|
||||
|
||||
type Service struct {
|
||||
httpClient *req.Client
|
||||
db *gorm.DB
|
||||
uploadManager *oss.UploaderManager
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
wsService *service.WebsocketService
|
||||
clientIds map[uint]string
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService, userService *service.UserService) *Service {
|
||||
return &Service{
|
||||
httpClient: req.C().SetTimeout(time.Minute * 3),
|
||||
db: db,
|
||||
taskQueue: store.NewRedisQueue("Video_Task_Queue", redisCli),
|
||||
notifyQueue: store.NewRedisQueue("Video_Notify_Queue", redisCli),
|
||||
wsService: wsService,
|
||||
uploadManager: manager,
|
||||
clientIds: map[uint]string{},
|
||||
userService: userService,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) PushTask(task types.VideoTask) {
|
||||
logger.Infof("add a new Video task to the task list: %+v", task)
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
|
||||
func (s *Service) Run() {
|
||||
// 将数据库中未提交的人物加载到队列
|
||||
var jobs []model.VideoJob
|
||||
s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
|
||||
for _, v := range jobs {
|
||||
var task types.VideoTask
|
||||
err := utils.JsonDecode(v.TaskInfo, &task)
|
||||
if err != nil {
|
||||
logger.Errorf("decode task info with error: %v", err)
|
||||
continue
|
||||
}
|
||||
task.Id = v.Id
|
||||
s.PushTask(task)
|
||||
s.clientIds[v.Id] = task.ClientId
|
||||
}
|
||||
logger.Info("Starting Video job consumer...")
|
||||
go func() {
|
||||
for {
|
||||
var task types.VideoTask
|
||||
err := s.taskQueue.LPop(&task)
|
||||
if err != nil {
|
||||
logger.Errorf("taking task with error: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// translate prompt
|
||||
if utils.HasChinese(task.Prompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId)
|
||||
if err == nil {
|
||||
task.Prompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if task.ClientId != "" {
|
||||
s.clientIds[task.Id] = task.ClientId
|
||||
}
|
||||
|
||||
var r LumaRespVo
|
||||
r, err = s.LumaCreate(task)
|
||||
if err != nil {
|
||||
logger.Errorf("create task with error: %v", err)
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"err_msg": err.Error(),
|
||||
"progress": service.FailTaskProgress,
|
||||
"cover_url": "/images/failed.jpg",
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
}
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
|
||||
// 更新任务信息
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"task_id": r.Id,
|
||||
"channel": r.Channel,
|
||||
"prompt_ext": r.Prompt,
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
s.PushTask(task)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
type LumaRespVo struct {
|
||||
Id string `json:"id"`
|
||||
Prompt string `json:"prompt"`
|
||||
State string `json:"state"`
|
||||
QueueState interface{} `json:"queue_state"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
Video interface{} `json:"video"`
|
||||
VideoRaw interface{} `json:"video_raw"`
|
||||
Liked interface{} `json:"liked"`
|
||||
EstimateWaitSeconds interface{} `json:"estimate_wait_seconds"`
|
||||
Thumbnail interface{} `json:"thumbnail"`
|
||||
Channel string `json:"channel,omitempty"`
|
||||
}
|
||||
|
||||
func (s *Service) LumaCreate(task types.VideoTask) (LumaRespVo, error) {
|
||||
// 读取 API KEY
|
||||
var apiKey model.ApiKey
|
||||
session := s.db.Session(&gorm.Session{}).Where("type", "luma").Where("enabled", true)
|
||||
if task.Channel != "" {
|
||||
session = session.Where("api_url", task.Channel)
|
||||
}
|
||||
tx := session.Order("last_used_at DESC").First(&apiKey)
|
||||
if tx.Error != nil {
|
||||
return LumaRespVo{}, errors.New("no available API KEY for Luma")
|
||||
}
|
||||
|
||||
reqBody := map[string]interface{}{
|
||||
"user_prompt": task.Prompt,
|
||||
"expand_prompt": task.Params.PromptOptimize,
|
||||
"loop": task.Params.Loop,
|
||||
"image_url": task.Params.StartImgURL,
|
||||
"image_end_url": task.Params.EndImgURL,
|
||||
}
|
||||
var res LumaRespVo
|
||||
apiURL := fmt.Sprintf("%s/luma/generations", apiKey.ApiURL)
|
||||
logger.Debugf("API URL: %s, request body: %+v", apiURL, reqBody)
|
||||
r, err := req.C().R().
|
||||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetBody(reqBody).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("请求 API 出错:%v", err)
|
||||
}
|
||||
|
||||
if r.StatusCode != 200 && r.StatusCode != 201 {
|
||||
return LumaRespVo{}, fmt.Errorf("请求 API 出错:%d, %s", r.StatusCode, r.String())
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
err = json.Unmarshal(body, &res)
|
||||
if err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("解析API数据失败:%v, %s", err, string(body))
|
||||
}
|
||||
|
||||
// update the last_use_at for api key
|
||||
apiKey.LastUsedAt = time.Now().Unix()
|
||||
session.Updates(&apiKey)
|
||||
res.Channel = apiKey.ApiURL
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
logger.Info("Running Suno task notify checking ...")
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
logger.Debugf("Receive notify message: %+v", message)
|
||||
client := s.wsService.Clients.Get(message.ClientId)
|
||||
if client == nil {
|
||||
continue
|
||||
}
|
||||
utils.SendChannelMsg(client, types.ChLuma, message.Message)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) DownloadFiles() {
|
||||
go func() {
|
||||
var items []model.VideoJob
|
||||
for {
|
||||
res := s.db.Where("progress", 102).Find(&items)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, v := range items {
|
||||
if v.WaterURL == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("try download video: %s", v.WaterURL)
|
||||
videoURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(v.WaterURL, true)
|
||||
if err != nil {
|
||||
logger.Errorf("download video with error: %v", err)
|
||||
continue
|
||||
}
|
||||
logger.Infof("download video success: %s", videoURL)
|
||||
v.WaterURL = videoURL
|
||||
|
||||
if v.VideoURL != "" {
|
||||
logger.Infof("try download no water video: %s", v.VideoURL)
|
||||
videoURL, err = s.uploadManager.GetUploadHandler().PutUrlFile(v.VideoURL, true)
|
||||
if err != nil {
|
||||
logger.Errorf("download video with error: %v", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
logger.Infof("download no water video success: %s", videoURL)
|
||||
v.VideoURL = videoURL
|
||||
v.Progress = 100
|
||||
s.db.Updates(&v)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.Id], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished})
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 10)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// SyncTaskProgress 异步拉取任务
|
||||
func (s *Service) SyncTaskProgress() {
|
||||
go func() {
|
||||
var jobs []model.VideoJob
|
||||
for {
|
||||
res := s.db.Where("progress < ?", 100).Where("task_id <> ?", "").Find(&jobs)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, job := range jobs {
|
||||
task, err := s.QueryLumaTask(job.TaskId, job.Channel)
|
||||
if err != nil {
|
||||
logger.Errorf("query task with error: %v", err)
|
||||
// 更新任务信息
|
||||
s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress, // 102 表示资源未下载完成,
|
||||
"err_msg": err.Error(),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Debugf("task: %+v", task)
|
||||
if task.State == "completed" { // 更新任务信息
|
||||
data := map[string]interface{}{
|
||||
"progress": 102, // 102 表示资源未下载完成,
|
||||
"water_url": task.Video.Url,
|
||||
"raw_data": utils.JsonEncode(task),
|
||||
"prompt_ext": task.Prompt,
|
||||
"cover_url": task.Thumbnail.Url,
|
||||
}
|
||||
if task.Video.DownloadUrl != "" {
|
||||
data["video_url"] = task.Video.DownloadUrl
|
||||
}
|
||||
err = s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(data).Error
|
||||
if err != nil {
|
||||
logger.Errorf("更新数据库失败:%v", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// 找出失败的任务,并恢复其扣减算力
|
||||
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
|
||||
for _, job := range jobs {
|
||||
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
|
||||
Type: types.PowerRefund,
|
||||
Model: "luma",
|
||||
Remark: fmt.Sprintf("Luma 任务失败,退回算力。任务ID:%s,Err:%s", job.TaskId, job.ErrMsg),
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// 更新任务状态
|
||||
s.db.Model(&job).UpdateColumn("power", 0)
|
||||
}
|
||||
time.Sleep(time.Second * 10)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
type LumaTaskVo struct {
|
||||
Id string `json:"id"`
|
||||
Liked interface{} `json:"liked"`
|
||||
State string `json:"state"`
|
||||
Video struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
Thumbnail string `json:"thumbnail"`
|
||||
DownloadUrl string `json:"download_url"`
|
||||
} `json:"video"`
|
||||
Prompt string `json:"prompt"`
|
||||
UserId string `json:"user_id"`
|
||||
BatchId string `json:"batch_id"`
|
||||
Thumbnail struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"thumbnail"`
|
||||
VideoRaw struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"video_raw"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
LastFrame struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"last_frame"`
|
||||
}
|
||||
|
||||
func (s *Service) QueryLumaTask(taskId string, channel string) (LumaTaskVo, error) {
|
||||
// 读取 API KEY
|
||||
var apiKey model.ApiKey
|
||||
err := s.db.Session(&gorm.Session{}).Where("type", "luma").
|
||||
Where("api_url", channel).
|
||||
Where("enabled", true).
|
||||
Order("last_used_at DESC").First(&apiKey).Error
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, errors.New("no available API KEY for Luma")
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/luma/generations/%s", apiKey.ApiURL, taskId)
|
||||
var res LumaTaskVo
|
||||
r, err := req.C().R().SetHeader("Authorization", "Bearer "+apiKey.Value).Get(apiURL)
|
||||
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, fmt.Errorf("请求 API 失败:%v", err)
|
||||
}
|
||||
defer r.Body.Close()
|
||||
|
||||
if r.StatusCode != 200 {
|
||||
return LumaTaskVo{}, fmt.Errorf("API 返回失败:%v", r.String())
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
err = json.Unmarshal(body, &res)
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, fmt.Errorf("解析API数据失败:%v, %s", err, string(body))
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
663
api/service/video/video.go
Normal file
663
api/service/video/video.go
Normal file
@@ -0,0 +1,663 @@
|
||||
package video
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
logger2 "geekai/logger"
|
||||
"geekai/service"
|
||||
"geekai/service/oss"
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/go-redis/redis/v8"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var logger = logger2.GetLogger()
|
||||
|
||||
type Service struct {
|
||||
httpClient *req.Client
|
||||
db *gorm.DB
|
||||
uploadManager *oss.UploaderManager
|
||||
taskQueue *store.RedisQueue
|
||||
wsService *service.WebsocketService
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService, userService *service.UserService) *Service {
|
||||
return &Service{
|
||||
httpClient: req.C().SetTimeout(time.Minute * 3),
|
||||
db: db,
|
||||
taskQueue: store.NewRedisQueue("Video_Task_Queue", redisCli),
|
||||
wsService: wsService,
|
||||
uploadManager: manager,
|
||||
userService: userService,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) PushTask(task types.VideoTask) {
|
||||
logger.Infof("add a new Video task to the task list: %+v", task)
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
|
||||
func (s *Service) Run() {
|
||||
// 将数据库中未提交的任务加载到队列
|
||||
var jobs []model.VideoJob
|
||||
s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
|
||||
for _, v := range jobs {
|
||||
var task types.VideoTask
|
||||
err := utils.JsonDecode(v.TaskInfo, &task)
|
||||
if err != nil {
|
||||
logger.Errorf("decode task info with error: %v", err)
|
||||
continue
|
||||
}
|
||||
task.Id = v.Id
|
||||
s.PushTask(task)
|
||||
}
|
||||
logger.Info("Starting Video job consumer...")
|
||||
go func() {
|
||||
for {
|
||||
var task types.VideoTask
|
||||
err := s.taskQueue.LPop(&task)
|
||||
if err != nil {
|
||||
logger.Errorf("taking task with error: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if task.Type == types.VideoLuma {
|
||||
// translate prompt
|
||||
if utils.HasChinese(task.Prompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId)
|
||||
if err == nil {
|
||||
task.Prompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
}
|
||||
}
|
||||
var r LumaRespVo
|
||||
r, err = s.LumaCreate(task)
|
||||
if err != nil {
|
||||
logger.Errorf("create task with error: %v", err)
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"err_msg": err.Error(),
|
||||
"progress": service.FailTaskProgress,
|
||||
"cover_url": "/images/failed.jpg",
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 更新任务信息
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"task_id": r.Id,
|
||||
"channel": r.Channel,
|
||||
"prompt_ext": r.Prompt,
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
s.PushTask(task)
|
||||
}
|
||||
} else if task.Type == types.VideoKeLing {
|
||||
var r KeLingRespVo
|
||||
r, err = s.KeLingCreate(task)
|
||||
logger.Debugf("ke ling create task result: %+v", r)
|
||||
|
||||
if err != nil {
|
||||
logger.Errorf("create task with error: %v", err)
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"err_msg": err.Error(),
|
||||
"progress": service.FailTaskProgress,
|
||||
"cover_url": "/images/failed.jpg",
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 更新任务信息
|
||||
err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||||
"task_id": r.Data.TaskID,
|
||||
"channel": r.Channel,
|
||||
"prompt_ext": task.Prompt,
|
||||
}).Error
|
||||
if err != nil {
|
||||
logger.Errorf("update task with error: %v", err)
|
||||
s.PushTask(task)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) DownloadFiles() {
|
||||
go func() {
|
||||
var items []model.VideoJob
|
||||
for {
|
||||
res := s.db.Where("progress", 102).Find(&items)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, v := range items {
|
||||
if v.WaterURL == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("try download video: %s", v.WaterURL)
|
||||
videoURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(v.WaterURL, true)
|
||||
if err != nil {
|
||||
logger.Errorf("download video with error: %v", err)
|
||||
continue
|
||||
}
|
||||
logger.Infof("download video success: %s", videoURL)
|
||||
v.WaterURL = videoURL
|
||||
|
||||
if v.VideoURL != "" {
|
||||
logger.Infof("try download no water video: %s", v.VideoURL)
|
||||
videoURL, err = s.uploadManager.GetUploadHandler().PutUrlFile(v.VideoURL, true)
|
||||
if err != nil {
|
||||
logger.Errorf("download video with error: %v", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
logger.Infof("download no water video success: %s", videoURL)
|
||||
v.VideoURL = videoURL
|
||||
v.Progress = 100
|
||||
s.db.Updates(&v)
|
||||
|
||||
// Convert TaskInfo to VideoTask
|
||||
var videoTask types.VideoTask
|
||||
if err := json.Unmarshal([]byte(v.TaskInfo), &videoTask); err != nil {
|
||||
logger.Errorf("failed to unmarshal task info to VideoTask: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 10)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// SyncTaskProgress 异步拉取任务
|
||||
func (s *Service) SyncTaskProgress() {
|
||||
go func() {
|
||||
var jobs []model.VideoJob
|
||||
for {
|
||||
res := s.db.Where("progress < ?", 100).Where("task_id <> ?", "").Find(&jobs)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, job := range jobs {
|
||||
if job.Type == types.VideoLuma {
|
||||
task, err := s.QueryLumaTask(job.TaskId, job.Channel)
|
||||
if err != nil {
|
||||
logger.Errorf("query task with error: %v", err)
|
||||
// 更新任务信息
|
||||
s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress, // 102 表示资源未下载完成,
|
||||
"err_msg": err.Error(),
|
||||
"cover_url": "/images/failed.jpg",
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Debugf("task: %+v", task)
|
||||
if task.State == "completed" { // 更新任务信息
|
||||
data := map[string]interface{}{
|
||||
"progress": 102, // 102 表示资源未下载完成,
|
||||
"water_url": task.Video.Url,
|
||||
"raw_data": utils.JsonEncode(task),
|
||||
"prompt_ext": task.Prompt,
|
||||
"cover_url": task.Thumbnail.Url,
|
||||
}
|
||||
if task.Video.DownloadUrl != "" {
|
||||
data["video_url"] = task.Video.DownloadUrl
|
||||
}
|
||||
err = s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(data).Error
|
||||
if err != nil {
|
||||
logger.Errorf("更新数据库失败:%v", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
} else if job.Type == types.VideoKeLing {
|
||||
// Convert TaskInfo to VideoTask
|
||||
var videoTask types.VideoTask
|
||||
if err := json.Unmarshal([]byte(job.TaskInfo), &videoTask); err != nil {
|
||||
logger.Errorf("failed to unmarshal task info to VideoTask: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Type assert task.Params to KeLingVideoParams
|
||||
paramsMap, ok := videoTask.Params.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert map to KeLingVideoParams
|
||||
paramsBytes, err := json.Marshal(paramsMap)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
var params types.KeLingVideoParams
|
||||
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
task, err := s.QueryKeLingTask(job.TaskId, job.Channel, params.TaskType)
|
||||
if err != nil {
|
||||
logger.Errorf("query task with error: %v", err)
|
||||
// 更新任务信息
|
||||
s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress, // 102 表示资源未下载完成,
|
||||
"err_msg": err.Error(),
|
||||
"cover_url": "/images/failed.jpg",
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Debugf("task: %+v", task)
|
||||
if task.TaskStatus == "succeed" { // 更新任务信息
|
||||
data := map[string]interface{}{
|
||||
"progress": 102, // 102 表示资源未下载完成,
|
||||
"water_url": task.TaskResult.Videos[0].URL,
|
||||
"raw_data": utils.JsonEncode(task),
|
||||
"prompt_ext": job.Prompt,
|
||||
"cover_url": "",
|
||||
}
|
||||
if len(task.TaskResult.Videos) > 0 {
|
||||
data["video_url"] = task.TaskResult.Videos[0].URL
|
||||
}
|
||||
err = s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(data).Error
|
||||
if err != nil {
|
||||
logger.Errorf("更新数据库失败:%v", err)
|
||||
continue
|
||||
}
|
||||
} else if task.TaskStatus == "failed" {
|
||||
// 更新任务信息
|
||||
s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress,
|
||||
"err_msg": task.TaskStatusMsg,
|
||||
"cover_url": "/images/failed.jpg",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// 找出失败的任务,并恢复其扣减算力
|
||||
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
|
||||
for _, job := range jobs {
|
||||
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
|
||||
Type: types.PowerRefund,
|
||||
Model: job.Type,
|
||||
Remark: fmt.Sprintf("%s 任务失败,退回算力。任务ID:%s,Err:%s", job.Type, job.TaskId, job.ErrMsg),
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// 更新任务状态
|
||||
s.db.Model(&job).UpdateColumn("power", 0)
|
||||
}
|
||||
time.Sleep(time.Second * 10)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
type LumaTaskVo struct {
|
||||
Id string `json:"id"`
|
||||
Liked interface{} `json:"liked"`
|
||||
State string `json:"state"`
|
||||
Video struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
Thumbnail string `json:"thumbnail"`
|
||||
DownloadUrl string `json:"download_url"`
|
||||
} `json:"video"`
|
||||
Prompt string `json:"prompt"`
|
||||
UserId string `json:"user_id"`
|
||||
BatchId string `json:"batch_id"`
|
||||
Thumbnail struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"thumbnail"`
|
||||
VideoRaw struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"video_raw"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
LastFrame struct {
|
||||
Url string `json:"url"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
} `json:"last_frame"`
|
||||
}
|
||||
|
||||
type LumaRespVo struct {
|
||||
Id string `json:"id"`
|
||||
Prompt string `json:"prompt"`
|
||||
State string `json:"state"`
|
||||
QueueState interface{} `json:"queue_state"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
Video interface{} `json:"video"`
|
||||
VideoRaw interface{} `json:"video_raw"`
|
||||
Liked interface{} `json:"liked"`
|
||||
EstimateWaitSeconds interface{} `json:"estimate_wait_seconds"`
|
||||
Thumbnail interface{} `json:"thumbnail"`
|
||||
Channel string `json:"channel,omitempty"`
|
||||
}
|
||||
|
||||
func (s *Service) LumaCreate(task types.VideoTask) (LumaRespVo, error) {
|
||||
// 读取 API KEY
|
||||
var apiKey model.ApiKey
|
||||
session := s.db.Session(&gorm.Session{}).Where("type", "luma").Where("enabled", true)
|
||||
if task.Channel != "" {
|
||||
session = session.Where("api_url", task.Channel)
|
||||
}
|
||||
tx := session.Order("last_used_at DESC").First(&apiKey)
|
||||
if tx.Error != nil {
|
||||
return LumaRespVo{}, errors.New("no available API KEY for Luma")
|
||||
}
|
||||
|
||||
// Type assert task.Params to LumaVideoParams
|
||||
paramsMap, ok := task.Params.(map[string]interface{})
|
||||
if !ok {
|
||||
return LumaRespVo{}, errors.New("invalid params type for Luma video task")
|
||||
}
|
||||
|
||||
// Convert map to LumaVideoParams
|
||||
paramsBytes, err := json.Marshal(paramsMap)
|
||||
if err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("failed to marshal params: %v", err)
|
||||
}
|
||||
|
||||
var params types.LumaVideoParams
|
||||
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("failed to unmarshal params: %v", err)
|
||||
}
|
||||
|
||||
reqBody := map[string]interface{}{
|
||||
"user_prompt": task.Prompt,
|
||||
"expand_prompt": params.PromptOptimize,
|
||||
"loop": params.Loop,
|
||||
"image_url": params.StartImgURL, // 图生视频
|
||||
"image_end_url": params.EndImgURL, // 图生视频
|
||||
}
|
||||
|
||||
var res LumaRespVo
|
||||
apiURL := fmt.Sprintf("%s/luma/generations", apiKey.ApiURL)
|
||||
logger.Debugf("API URL: %s, request body: %+v", apiURL, reqBody)
|
||||
r, err := req.C().R().
|
||||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetBody(reqBody).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("请求 API 出错:%v", err)
|
||||
}
|
||||
|
||||
if r.StatusCode != 200 && r.StatusCode != 201 {
|
||||
return LumaRespVo{}, fmt.Errorf("请求 API 出错:%d, %s", r.StatusCode, r.String())
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
err = json.Unmarshal(body, &res)
|
||||
if err != nil {
|
||||
return LumaRespVo{}, fmt.Errorf("解析API数据失败:%v, %s", err, string(body))
|
||||
}
|
||||
|
||||
// update the last_use_at for api key
|
||||
apiKey.LastUsedAt = time.Now().Unix()
|
||||
session.Updates(&apiKey)
|
||||
res.Channel = apiKey.ApiURL
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Service) QueryLumaTask(taskId string, channel string) (LumaTaskVo, error) {
|
||||
// 读取 API KEY
|
||||
var apiKey model.ApiKey
|
||||
err := s.db.Session(&gorm.Session{}).Where("type", "luma").
|
||||
Where("api_url", channel).
|
||||
Where("enabled", true).
|
||||
Order("last_used_at DESC").First(&apiKey).Error
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, errors.New("no available API KEY for Luma")
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/luma/generations/%s", apiKey.ApiURL, taskId)
|
||||
var res LumaTaskVo
|
||||
r, err := req.C().R().SetHeader("Authorization", "Bearer "+apiKey.Value).Get(apiURL)
|
||||
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, fmt.Errorf("请求 API 失败:%v", err)
|
||||
}
|
||||
defer r.Body.Close()
|
||||
|
||||
if r.StatusCode != 200 {
|
||||
return LumaTaskVo{}, fmt.Errorf("API 返回失败:%v", r.String())
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
err = json.Unmarshal(body, &res)
|
||||
if err != nil {
|
||||
return LumaTaskVo{}, fmt.Errorf("解析API数据失败:%v, %s", err, string(body))
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
type KeLingRespVo struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
RequestID string `json:"request_id"`
|
||||
Data struct {
|
||||
TaskID string `json:"task_id"`
|
||||
TaskStatus string `json:"task_status"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
UpdatedAt int64 `json:"updated_at"`
|
||||
} `json:"data"`
|
||||
Channel string `json:"channel,omitempty"`
|
||||
}
|
||||
|
||||
func (s *Service) KeLingCreate(task types.VideoTask) (KeLingRespVo, error) {
|
||||
var apiKey model.ApiKey
|
||||
session := s.db.Session(&gorm.Session{}).Where("type", "keling").Where("enabled", true)
|
||||
if task.Channel != "" {
|
||||
session = session.Where("api_url", task.Channel)
|
||||
}
|
||||
tx := session.Order("last_used_at DESC").First(&apiKey)
|
||||
if tx.Error != nil {
|
||||
return KeLingRespVo{}, errors.New("no available API KEY for keling")
|
||||
}
|
||||
|
||||
// Type assert task.Params to KeLingVideoParams
|
||||
paramsMap, ok := task.Params.(map[string]interface{})
|
||||
if !ok {
|
||||
return KeLingRespVo{}, errors.New("invalid params type for KeLing video task")
|
||||
}
|
||||
|
||||
// Convert map to KeLingVideoParams
|
||||
paramsBytes, err := json.Marshal(paramsMap)
|
||||
if err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to marshal params: %v", err)
|
||||
}
|
||||
|
||||
var params types.KeLingVideoParams
|
||||
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to unmarshal params: %v", err)
|
||||
}
|
||||
|
||||
// 2. 构建API请求参数
|
||||
payload := map[string]interface{}{
|
||||
"model_name": params.Model,
|
||||
"prompt": task.Prompt,
|
||||
"negative_prompt": params.NegPrompt,
|
||||
"cfg_scale": params.CfgScale,
|
||||
"mode": params.Mode,
|
||||
"aspect_ratio": params.AspectRatio,
|
||||
"duration": params.Duration,
|
||||
}
|
||||
|
||||
// 只有当 CameraControl 的类型不为空时,才处理摄像机控制参数
|
||||
if params.CameraControl.Type != "" {
|
||||
cameraControl := map[string]interface{}{
|
||||
"type": params.CameraControl.Type,
|
||||
}
|
||||
|
||||
// 只有在 simple 类型时才添加 config 参数
|
||||
if params.CameraControl.Type == "simple" {
|
||||
cameraControl["config"] = params.CameraControl.Config
|
||||
}
|
||||
|
||||
payload["camera_control"] = cameraControl
|
||||
}
|
||||
|
||||
// 处理图生视频
|
||||
if params.TaskType == "image2video" {
|
||||
payload["image"] = params.Image
|
||||
payload["image_tail"] = params.ImageTail
|
||||
}
|
||||
|
||||
jsonPayload, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to marshal payload: %v", err)
|
||||
}
|
||||
|
||||
// 3. 准备HTTP请求
|
||||
url := fmt.Sprintf("%s/kling/v1/videos/%s", apiKey.ApiURL, params.TaskType)
|
||||
req, err := http.NewRequest("POST", url, bytes.NewReader(jsonPayload))
|
||||
if err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to create request: %v", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey.Value)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// 4. 发送请求
|
||||
client := &http.Client{Timeout: time.Duration(30) * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to send request: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// 5. 处理响应
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to read response: %v", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return KeLingRespVo{}, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var apiResponse = KeLingRespVo{}
|
||||
if err := json.Unmarshal(body, &apiResponse); err != nil {
|
||||
return KeLingRespVo{}, fmt.Errorf("failed to parse response: %v", err)
|
||||
}
|
||||
// 设置 API 通道
|
||||
apiResponse.Channel = apiKey.ApiURL
|
||||
return apiResponse, nil
|
||||
}
|
||||
|
||||
// VideoCallbackData 表示视频生成任务的回调数据
|
||||
type VideoCallbackData struct {
|
||||
TaskID string `json:"task_id"`
|
||||
TaskStatus string `json:"task_status"`
|
||||
TaskStatusMsg string `json:"task_status_msg"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
UpdatedAt int64 `json:"updated_at"`
|
||||
TaskResult TaskResult `json:"task_result"`
|
||||
}
|
||||
|
||||
type TaskResult struct {
|
||||
Images []CallBackImageResult `json:"images,omitempty"`
|
||||
Videos []CallBackVideoResult `json:"videos,omitempty"`
|
||||
}
|
||||
|
||||
type CallBackImageResult struct {
|
||||
Index int `json:"index"`
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
type CallBackVideoResult struct {
|
||||
ID string `json:"id"`
|
||||
URL string `json:"url"`
|
||||
Duration string `json:"duration"`
|
||||
}
|
||||
|
||||
func (s *Service) QueryKeLingTask(taskId string, channel string, action string) (VideoCallbackData, error) {
|
||||
var apiKey model.ApiKey
|
||||
err := s.db.Session(&gorm.Session{}).Where("type", "keling").
|
||||
//Where("api_url", channel).
|
||||
Where("enabled", true).
|
||||
Order("last_used_at DESC").First(&apiKey).Error
|
||||
if err != nil {
|
||||
return VideoCallbackData{}, errors.New("no available API KEY for keling")
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s/kling/v1/videos/%s/%s", apiKey.ApiURL, action, taskId)
|
||||
req, err := http.NewRequest("GET", url, nil)
|
||||
if err != nil {
|
||||
return VideoCallbackData{}, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey.Value)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := &http.Client{}
|
||||
res, err := client.Do(req)
|
||||
if err != nil {
|
||||
return VideoCallbackData{}, fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer res.Body.Close()
|
||||
|
||||
if res.StatusCode != http.StatusOK {
|
||||
return VideoCallbackData{}, fmt.Errorf("unexpected status code: %d", res.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return VideoCallbackData{}, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
var response struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data VideoCallbackData `json:"data"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(body, &response); err != nil {
|
||||
return VideoCallbackData{}, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
if response.Code != 0 {
|
||||
return VideoCallbackData{}, fmt.Errorf("API error: %s", response.Message)
|
||||
}
|
||||
|
||||
return response.Data, nil
|
||||
}
|
||||
Reference in New Issue
Block a user