From 135755d21d307af024e19a6d3f693772e0c572a0 Mon Sep 17 00:00:00 2001 From: RockYang Date: Fri, 8 Nov 2024 18:06:39 +0800 Subject: [PATCH] enable to set the translate model --- api/core/types/config.go | 6 +-- api/core/types/task.go | 63 +++++++++++++------------ api/handler/chat_handler.go | 2 +- api/handler/chat_model_handler.go | 20 ++++---- api/handler/dalle_handler.go | 17 +++---- api/handler/function_handler.go | 7 ++- api/handler/markmap_handler.go | 2 +- api/handler/mj_handler.go | 49 ++++++++++---------- api/handler/prompt_handler.go | 59 ++++++++++++++++++++++++ api/handler/sd_handler.go | 53 +++++++++++---------- api/handler/suno_handler.go | 37 --------------- api/handler/user_handler.go | 34 +++++++------- api/handler/video_handler.go | 13 +++--- api/main.go | 6 ++- api/service/dalle/service.go | 2 +- api/service/mj/service.go | 5 +- api/service/sd/service.go | 19 ++++++-- api/service/types.go | 76 ++++++++++++++++++++++++++++++- api/service/video/luma.go | 2 +- api/store/model/sd_job.go | 1 + api/store/vo/mj_job.go | 33 +++++++------- api/utils/openai.go | 19 +++++--- database/update-v4.1.7.sql | 1 + web/src/views/Suno.vue | 2 +- web/src/views/admin/SysConfig.vue | 7 ++- 25 files changed, 328 insertions(+), 207 deletions(-) create mode 100644 api/handler/prompt_handler.go create mode 100644 database/update-v4.1.7.sql diff --git a/api/core/types/config.go b/api/core/types/config.go index b827d1f3..a033c17c 100644 --- a/api/core/types/config.go +++ b/api/core/types/config.go @@ -142,7 +142,6 @@ type SystemConfig struct { OrderPayTimeout int `json:"order_pay_timeout,omitempty"` //订单支付超时时间 VipInfoText string `json:"vip_info_text,omitempty"` // 会员页面充值说明 - DefaultModels []int `json:"default_models,omitempty"` // 默认开通的 AI 模型 MjPower int `json:"mj_power,omitempty"` // MJ 绘画消耗算力 MjActionPower int `json:"mj_action_power,omitempty"` // MJ 操作（放大，变换）消耗算力 @@ -164,6 +163,7 @@ type SystemConfig struct { Copyright string `json:"copyright"` // 版权信息 MarkMapText string `json:"mark_map_text"` // 思维导入的默认文本 - EnabledVerify bool `json:"enabled_verify"` // 是否启用验证码 - EmailWhiteList []string `json:"email_white_list"` // 邮箱白名单列表 + EnabledVerify bool `json:"enabled_verify"` // 是否启用验证码 + EmailWhiteList []string `json:"email_white_list"` // 邮箱白名单列表 + TranslateModelId int `json:"translate_model_id"` // 用来做提示词翻译的大模型 id } diff --git a/api/core/types/task.go b/api/core/types/task.go index 5dac6443..444e25c0 100644 --- a/api/core/types/task.go +++ b/api/core/types/task.go @@ -24,30 +24,31 @@ const ( // MjTask MidJourney 任务 type MjTask struct { - Id uint `json:"id"` // 任务ID - TaskId string `json:"task_id"` // 中转任务ID - ClientId string `json:"client_id"` - ImgArr []string `json:"img_arr"` - Type TaskType `json:"type"` - UserId int `json:"user_id"` - Prompt string `json:"prompt,omitempty"` - NegPrompt string `json:"neg_prompt,omitempty"` - Params string `json:"full_prompt"` - Index int `json:"index,omitempty"` - MessageId string `json:"message_id,omitempty"` - MessageHash string `json:"message_hash,omitempty"` - RetryCount int `json:"retry_count"` - ChannelId string `json:"channel_id"` // 渠道ID，用来区分是哪个渠道创建的任务，一个任务的 create 和 action 操作必须要再同一个渠道 - Mode string `json:"mode"` // 绘画模式，relax, fast, turbo + Id uint `json:"id"` // 任务ID + TaskId string `json:"task_id"` // 中转任务ID + ClientId string `json:"client_id"` + ImgArr []string `json:"img_arr"` + Type TaskType `json:"type"` + UserId int `json:"user_id"` + Prompt string `json:"prompt,omitempty"` + NegPrompt string `json:"neg_prompt,omitempty"` + Params string `json:"full_prompt"` + Index int `json:"index,omitempty"` + MessageId string `json:"message_id,omitempty"` + MessageHash string `json:"message_hash,omitempty"` + ChannelId string `json:"channel_id"` // 渠道ID，用来区分是哪个渠道创建的任务，一个任务的 create 和 action 操作必须要再同一个渠道 + Mode string `json:"mode"` // 绘画模式，relax, fast, turbo + TranslateModelId int `json:"translate_model_id"` // 提示词翻译模型ID } type SdTask struct { - Id int `json:"id"` // job 数据库ID - Type TaskType `json:"type"` - ClientId string `json:"client_id"` - UserId int `json:"user_id"` - Params SdTaskParams `json:"params"` - RetryCount int `json:"retry_count"` + Id int `json:"id"` // job 数据库ID + Type TaskType `json:"type"` + ClientId string `json:"client_id"` + UserId int `json:"user_id"` + Params SdTaskParams `json:"params"` + RetryCount int `json:"retry_count"` + TranslateModelId int `json:"translate_model_id"` // 提示词翻译模型ID } type SdTaskParams struct { @@ -81,7 +82,8 @@ type DallTask struct { Size string `json:"size"` Style string `json:"style"` - Power int `json:"power"` + Power int `json:"power"` + TranslateModelId int `json:"translate_model_id"` // 提示词翻译模型ID } type SunoTask struct { @@ -109,14 +111,15 @@ const ( ) type VideoTask struct { - ClientId string `json:"client_id"` - Id uint `json:"id"` - Channel string `json:"channel"` - UserId int `json:"user_id"` - Type string `json:"type"` - TaskId string `json:"task_id"` - Prompt string `json:"prompt"` // 提示词 - Params VideoParams `json:"params"` + ClientId string `json:"client_id"` + Id uint `json:"id"` + Channel string `json:"channel"` + UserId int `json:"user_id"` + Type string `json:"type"` + TaskId string `json:"task_id"` + Prompt string `json:"prompt"` // 提示词 + Params VideoParams `json:"params"` + TranslateModelId int `json:"translate_model_id"` // 提示词翻译模型ID } type VideoParams struct { diff --git a/api/handler/chat_handler.go b/api/handler/chat_handler.go index 4edffa78..f61a7c53 100644 --- a/api/handler/chat_handler.go +++ b/api/handler/chat_handler.go @@ -371,7 +371,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, sessi } else { client = http.DefaultClient } - logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model) + logger.Infof("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model) request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey.Value)) // 更新API KEY 最后使用时间 h.DB.Model(&model.ApiKey{}).Where("id", apiKey.Id).UpdateColumn("last_used_at", time.Now().Unix()) diff --git a/api/handler/chat_model_handler.go b/api/handler/chat_model_handler.go index 1b74f348..d054ef3c 100644 --- a/api/handler/chat_model_handler.go +++ b/api/handler/chat_model_handler.go @@ -30,29 +30,25 @@ func NewChatModelHandler(app *core.AppServer, db *gorm.DB) *ChatModelHandler { func (h *ChatModelHandler) List(c *gin.Context) { var items []model.ChatModel var chatModels = make([]vo.ChatModel, 0) - var res *gorm.DB session := h.DB.Session(&gorm.Session{}).Where("enabled", true) t := c.Query("type") if t != "" { session = session.Where("type", t) } - // 如果用户没有登录，则加载所有开放模型 - if !h.IsLogin(c) { - res = session.Where("open", true).Order("sort_num ASC").Find(&items) - } else { + + session = session.Where("open", true) + if h.IsLogin(c) { user, _ := h.GetLoginUser(c) var models []int err := utils.JsonDecode(user.ChatModels, &models) - if err != nil { - resp.ERROR(c, "当前用户没有订阅任何模型") - return - } // 查询用户有权限访问的模型以及所有开放的模型 - res = h.DB.Where("enabled = ?", true).Where( - h.DB.Where("id IN ?", models).Or("open", true), - ).Order("sort_num ASC").Find(&items) + if err == nil { + session = session.Or("id IN ?", models) + } + } + res := session.Order("sort_num ASC").Find(&items) if res.Error == nil { for _, item := range items { var cm vo.ChatModel diff --git a/api/handler/dalle_handler.go b/api/handler/dalle_handler.go index 404c9704..8636b167 100644 --- a/api/handler/dalle_handler.go +++ b/api/handler/dalle_handler.go @@ -84,14 +84,15 @@ func (h *DallJobHandler) Image(c *gin.Context) { } h.dallService.PushTask(types.DallTask{ - ClientId: data.ClientId, - JobId: job.Id, - UserId: uint(userId), - Prompt: data.Prompt, - Quality: data.Quality, - Size: data.Size, - Style: data.Style, - Power: job.Power, + ClientId: data.ClientId, + JobId: job.Id, + UserId: uint(userId), + Prompt: data.Prompt, + Quality: data.Quality, + Size: data.Size, + Style: data.Style, + Power: job.Power, + TranslateModelId: h.App.SysConfig.TranslateModelId, }) resp.SUCCESS(c) } diff --git a/api/handler/function_handler.go b/api/handler/function_handler.go index f1838d4d..b39ea6ff 100644 --- a/api/handler/function_handler.go +++ b/api/handler/function_handler.go @@ -113,10 +113,13 @@ func (h *FunctionHandler) WeiBo(c *gin.Context) { SetHeader("AppId", h.config.AppId). SetHeader("Authorization", fmt.Sprintf("Bearer %s", h.config.Token)). SetSuccessResult(&res).Get(url) - if err != nil || r.IsErrorState() { - resp.ERROR(c, fmt.Sprintf("%v%v", err, r.Err)) + if err != nil { + resp.ERROR(c, fmt.Sprintf("%v", err)) return } + if r.IsErrorState() { + resp.ERROR(c, fmt.Sprintf("error http code status: %v", r.Status)) + } if res.Code != types.Success { resp.ERROR(c, res.Message) diff --git a/api/handler/markmap_handler.go b/api/handler/markmap_handler.go index e57f6f4d..3d7a7b14 100644 --- a/api/handler/markmap_handler.go +++ b/api/handler/markmap_handler.go @@ -87,7 +87,7 @@ func (h *MarkMapHandler) Generate(c *gin.Context) { 请直接生成结果，不要任何解释性语句。 `}) messages = append(messages, types.Message{Role: "user", Content: fmt.Sprintf("请生成一份有关【%s】一份思维导图，要求结构清晰，有条理", data.Prompt)}) - content, err := utils.SendOpenAIMessage(h.DB, messages, chatModel.Value, chatModel.KeyId) + content, err := utils.SendOpenAIMessage(h.DB, messages, data.ModelId) if err != nil { resp.ERROR(c, fmt.Sprintf("请求 OpenAI API 失败: %s", err)) return diff --git a/api/handler/mj_handler.go b/api/handler/mj_handler.go index 858a0d89..9edacca9 100644 --- a/api/handler/mj_handler.go +++ b/api/handler/mj_handler.go @@ -176,16 +176,17 @@ func (h *MidJourneyHandler) Image(c *gin.Context) { } h.mjService.PushTask(types.MjTask{ - Id: job.Id, - ClientId: data.ClientId, - TaskId: taskId, - Type: types.TaskType(data.TaskType), - Prompt: data.Prompt, - NegPrompt: data.NegPrompt, - Params: params, - UserId: userId, - ImgArr: data.ImgArr, - Mode: h.App.SysConfig.MjMode, + Id: job.Id, + ClientId: data.ClientId, + TaskId: taskId, + Type: types.TaskType(data.TaskType), + Prompt: data.Prompt, + NegPrompt: data.NegPrompt, + Params: params, + UserId: userId, + ImgArr: data.ImgArr, + Mode: h.App.SysConfig.MjMode, + TranslateModelId: h.App.SysConfig.TranslateModelId, }) // update user's power @@ -226,13 +227,12 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) { userId := utils.IntValue(utils.InterfaceToString(idValue), 0) taskId, _ := h.snowflake.Next(true) job := model.MidJourneyJob{ - Type: types.TaskUpscale.String(), - ReferenceId: data.MessageId, - UserId: userId, - TaskId: taskId, - Progress: 0, - Power: h.App.SysConfig.MjActionPower, - CreatedAt: time.Now(), + Type: types.TaskUpscale.String(), + UserId: userId, + TaskId: taskId, + Progress: 0, + Power: h.App.SysConfig.MjActionPower, + CreatedAt: time.Now(), } if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 { resp.ERROR(c, "添加任务失败："+res.Error.Error()) @@ -281,14 +281,13 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) { userId := utils.IntValue(utils.InterfaceToString(idValue), 0) taskId, _ := h.snowflake.Next(true) job := model.MidJourneyJob{ - Type: types.TaskVariation.String(), - ChannelId: data.ChannelId, - ReferenceId: data.MessageId, - UserId: userId, - TaskId: taskId, - Progress: 0, - Power: h.App.SysConfig.MjActionPower, - CreatedAt: time.Now(), + Type: types.TaskVariation.String(), + ChannelId: data.ChannelId, + UserId: userId, + TaskId: taskId, + Progress: 0, + Power: h.App.SysConfig.MjActionPower, + CreatedAt: time.Now(), } if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 { resp.ERROR(c, "添加任务失败："+res.Error.Error()) diff --git a/api/handler/prompt_handler.go b/api/handler/prompt_handler.go new file mode 100644 index 00000000..7e719803 --- /dev/null +++ b/api/handler/prompt_handler.go @@ -0,0 +1,59 @@ +package handler + +// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ +// * 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 ( + "fmt" + "geekai/core" + "geekai/core/types" + "geekai/service" + "geekai/service/oss" + "geekai/service/suno" + "geekai/utils" + "geekai/utils/resp" + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// 提示词生成 handler +// 使用 AI 生成绘画指令，歌词，视频生成指令等 + +type PromptHandler struct { + BaseHandler + sunoService *suno.Service + uploader *oss.UploaderManager + userService *service.UserService +} + +func NewPromptHandler(app *core.AppServer, db *gorm.DB, userService *service.UserService) *PromptHandler { + return &PromptHandler{ + BaseHandler: BaseHandler{ + App: app, + DB: db, + }, + userService: userService, + } +} + +// Lyric 生成歌词 +func (h *PromptHandler) Lyric(c *gin.Context) { + var data struct { + Prompt string `json:"prompt"` + } + if err := c.ShouldBindJSON(&data); err != nil { + resp.ERROR(c, types.InvalidArgs) + return + } + content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(service.LyricPromptTemplate, data.Prompt), h.App.SysConfig.TranslateModelId) + if err != nil { + resp.ERROR(c, err.Error()) + return + } + + resp.SUCCESS(c, content) +} diff --git a/api/handler/sd_handler.go b/api/handler/sd_handler.go index 437dceac..7c41cbb1 100644 --- a/api/handler/sd_handler.go +++ b/api/handler/sd_handler.go @@ -109,29 +109,37 @@ func (h *SdJobHandler) Image(c *gin.Context) { resp.ERROR(c, "error with generate task id: "+err.Error()) return } - params := types.SdTaskParams{ - TaskId: taskId, - Prompt: data.Prompt, - NegPrompt: data.NegPrompt, - Steps: data.Steps, - Sampler: data.Sampler, - FaceFix: data.FaceFix, - CfgScale: data.CfgScale, - Seed: data.Seed, - Height: data.Height, - Width: data.Width, - HdFix: data.HdFix, - HdRedrawRate: data.HdRedrawRate, - HdScale: data.HdScale, - HdScaleAlg: data.HdScaleAlg, - HdSteps: data.HdSteps, + + task := types.SdTask{ + ClientId: data.ClientId, + Type: types.TaskImage, + Params: types.SdTaskParams{ + TaskId: taskId, + Prompt: data.Prompt, + NegPrompt: data.NegPrompt, + Steps: data.Steps, + Sampler: data.Sampler, + FaceFix: data.FaceFix, + CfgScale: data.CfgScale, + Seed: data.Seed, + Height: data.Height, + Width: data.Width, + HdFix: data.HdFix, + HdRedrawRate: data.HdRedrawRate, + HdScale: data.HdScale, + HdScaleAlg: data.HdScaleAlg, + HdSteps: data.HdSteps, + }, + UserId: userId, + TranslateModelId: h.App.SysConfig.TranslateModelId, } job := model.SdJob{ UserId: userId, Type: types.TaskImage.String(), - TaskId: params.TaskId, - Params: utils.JsonEncode(params), + TaskId: taskId, + Params: utils.JsonEncode(task.Params), + TaskInfo: utils.JsonEncode(task), Prompt: data.Prompt, Progress: 0, Power: h.App.SysConfig.SdPower, @@ -143,13 +151,8 @@ func (h *SdJobHandler) Image(c *gin.Context) { return } - h.sdService.PushTask(types.SdTask{ - Id: int(job.Id), - ClientId: data.ClientId, - Type: types.TaskImage, - Params: params, - UserId: userId, - }) + task.Id = int(job.Id) + h.sdService.PushTask(task) // update user's power err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{ diff --git a/api/handler/suno_handler.go b/api/handler/suno_handler.go index d284c66f..a153c049 100644 --- a/api/handler/suno_handler.go +++ b/api/handler/suno_handler.go @@ -334,40 +334,3 @@ func (h *SunoHandler) Play(c *gin.Context) { } h.DB.Model(&model.SunoJob{}).Where("song_id", songId).UpdateColumn("play_times", gorm.Expr("play_times + ?", 1)) } - -const genLyricTemplate = ` -你是一位才华横溢的作曲家，拥有丰富的情感和细腻的笔触，你对文字有着独特的感悟力，能将各种情感和意境巧妙地融入歌词中。 -请以【%s】为主题创作一首歌曲，歌曲时间不要太短，3分钟左右，不要输出任何解释性的内容。 -输出格式如下： -歌曲名称 -第一节： -{{歌词内容}} -副歌： -{{歌词内容}} - -第二节： -{{歌词内容}} -副歌： -{{歌词内容}} - -尾声： -{{歌词内容}} -` - -// Lyric 生成歌词 -func (h *SunoHandler) Lyric(c *gin.Context) { - var data struct { - Prompt string `json:"prompt"` - } - if err := c.ShouldBindJSON(&data); err != nil { - resp.ERROR(c, types.InvalidArgs) - return - } - content, err := utils.OpenAIRequest(h.DB, fmt.Sprintf(genLyricTemplate, data.Prompt), "gpt-4o-mini", 0) - if err != nil { - resp.ERROR(c, err.Error()) - return - } - - resp.SUCCESS(c, content) -} diff --git a/api/handler/user_handler.go b/api/handler/user_handler.go index fd8a8046..3bc42ec1 100644 --- a/api/handler/user_handler.go +++ b/api/handler/user_handler.go @@ -132,14 +132,13 @@ func (h *UserHandler) Register(c *gin.Context) { salt := utils.RandString(8) user := model.User{ - Username: data.Username, - Password: utils.GenPassword(data.Password, salt), - Avatar: "/images/avatar/user.png", - Salt: salt, - Status: true, - ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色 - ChatModels: utils.JsonEncode(h.App.SysConfig.DefaultModels), // 默认开通的模型 - Power: h.App.SysConfig.InitPower, + Username: data.Username, + Password: utils.GenPassword(data.Password, salt), + Avatar: "/images/avatar/user.png", + Salt: salt, + Status: true, + ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色 + Power: h.App.SysConfig.InitPower, } // check if the username is existing @@ -417,16 +416,15 @@ func (h *UserHandler) CLoginCallback(c *gin.Context) { salt := utils.RandString(8) password := fmt.Sprintf("%d", utils.RandomNumber(8)) user = model.User{ - Username: fmt.Sprintf("%s@%d", loginType, utils.RandomNumber(10)), - Password: utils.GenPassword(password, salt), - Avatar: fmt.Sprintf("%s", data["avatar"]), - Salt: salt, - Status: true, - ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色 - ChatModels: utils.JsonEncode(h.App.SysConfig.DefaultModels), // 默认开通的模型 - Power: h.App.SysConfig.InitPower, - OpenId: fmt.Sprintf("%s", data["openid"]), - Nickname: fmt.Sprintf("%s", data["nickname"]), + Username: fmt.Sprintf("%s@%d", loginType, utils.RandomNumber(10)), + Password: utils.GenPassword(password, salt), + Avatar: fmt.Sprintf("%s", data["avatar"]), + Salt: salt, + Status: true, + ChatRoles: utils.JsonEncode([]string{"gpt"}), // 默认只订阅通用助手角色 + Power: h.App.SysConfig.InitPower, + OpenId: fmt.Sprintf("%s", data["openid"]), + Nickname: fmt.Sprintf("%s", data["nickname"]), } tx = h.DB.Create(&user) diff --git a/api/handler/video_handler.go b/api/handler/video_handler.go index aaa0bd86..42731164 100644 --- a/api/handler/video_handler.go +++ b/api/handler/video_handler.go @@ -96,12 +96,13 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) { // 创建任务 h.videoService.PushTask(types.VideoTask{ - ClientId: data.ClientId, - Id: job.Id, - UserId: userId, - Type: types.VideoLuma, - Prompt: data.Prompt, - Params: params, + ClientId: data.ClientId, + Id: job.Id, + UserId: userId, + Type: types.VideoLuma, + Prompt: data.Prompt, + Params: params, + TranslateModelId: h.App.SysConfig.TranslateModelId, }) // update user's power diff --git a/api/main.go b/api/main.go index 8fd36c82..17fb87ba 100644 --- a/api/main.go +++ b/api/main.go @@ -484,7 +484,6 @@ func main() { group.POST("update", h.Update) group.GET("detail", h.Detail) group.GET("play", h.Play) - group.POST("lyric", h.Lyric) }), fx.Provide(handler.NewVideoHandler), fx.Invoke(func(s *core.AppServer, h *handler.VideoHandler) { @@ -518,6 +517,11 @@ func main() { fx.Invoke(func(s *core.AppServer, h *handler.WebsocketHandler) { s.Engine.Any("/api/ws", h.Client) }), + fx.Provide(handler.NewPromptHandler), + fx.Invoke(func(s *core.AppServer, h *handler.PromptHandler) { + group := s.Engine.Group("/api/prompt") + group.POST("/lyric", h.Lyric) + }), fx.Invoke(func(s *core.AppServer, db *gorm.DB) { go func() { err := s.Run(db) diff --git a/api/service/dalle/service.go b/api/service/dalle/service.go index 12bef395..1d29a731 100644 --- a/api/service/dalle/service.go +++ b/api/service/dalle/service.go @@ -114,7 +114,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) { prompt := task.Prompt // translate prompt if utils.HasChinese(prompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, prompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, prompt), task.TranslateModelId) if err == nil { prompt = content logger.Debugf("重写后提示词：%s", prompt) diff --git a/api/service/mj/service.go b/api/service/mj/service.go index 72192dcf..3a22a51e 100644 --- a/api/service/mj/service.go +++ b/api/service/mj/service.go @@ -58,7 +58,7 @@ func (s *Service) Run() { // translate prompt if utils.HasChinese(task.Prompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId) if err == nil { task.Prompt = content } else { @@ -67,7 +67,7 @@ func (s *Service) Run() { } // translate negative prompt if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt), task.TranslateModelId) if err == nil { task.NegPrompt = content } else { @@ -275,7 +275,6 @@ func (s *Service) SyncTaskProgress() { } oldProgress := job.Progress job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0) - job.Prompt = task.PromptEn if task.ImageUrl != "" { job.OrgURL = task.ImageUrl } diff --git a/api/service/sd/service.go b/api/service/sd/service.go index 9bfd1ecd..fd196049 100644 --- a/api/service/sd/service.go +++ b/api/service/sd/service.go @@ -48,6 +48,19 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelD } func (s *Service) Run() { + // 将数据库中未提交的人物加载到队列 + var jobs []model.SdJob + s.db.Where("progress", 0).Find(&jobs) + for _, v := range jobs { + var task types.SdTask + err := utils.JsonDecode(v.TaskInfo, &task) + if err != nil { + logger.Errorf("decode task info with error: %v", err) + continue + } + task.Id = int(v.Id) + s.PushTask(task) + } logger.Infof("Starting Stable-Diffusion job consumer") go func() { for { @@ -60,7 +73,7 @@ func (s *Service) Run() { // translate prompt if utils.HasChinese(task.Params.Prompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Params.Prompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Params.Prompt), task.TranslateModelId) if err == nil { task.Params.Prompt = content } else { @@ -70,7 +83,7 @@ func (s *Service) Run() { // translate negative prompt if task.Params.NegPrompt != "" && utils.HasChinese(task.Params.NegPrompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Params.NegPrompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Params.NegPrompt), task.TranslateModelId) if err == nil { task.Params.NegPrompt = content } else { @@ -161,7 +174,7 @@ func (s *Service) Txt2Img(task types.SdTask) error { } apiURL := fmt.Sprintf("%s/sdapi/v1/txt2img", apiKey.ApiURL) - logger.Debugf("send image request to %s", apiURL) + logger.Infof("send image request to %s", apiURL) // send a request to sd api endpoint go func() { response, err := s.httpClient.R(). diff --git a/api/service/types.go b/api/service/types.go index 1c5c601e..ad9ef493 100644 --- a/api/service/types.go +++ b/api/service/types.go @@ -14,5 +14,79 @@ type NotifyMessage struct { Message string `json:"message"` } -const RewritePromptTemplate = "Please rewrite the following text into AI painting prompt words, and please try to add detailed description of the picture, painting style, scene, rendering effect, picture light and other creative elements. Just output the final prompt word directly. Do not output any explanation lines. The text to be rewritten is: [%s]" 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. + +Please specify details about the artwork, such as the style, subject, mood, and other important characteristics you want the resulting image to have. + +Remember, prompts should always be output in English. + +# 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. + +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】 +` + +const LyricPromptTemplate = ` +你是一位才华横溢的作曲家，拥有丰富的情感和细腻的笔触，你对文字有着独特的感悟力，能将各种情感和意境巧妙地融入歌词中。 +请以【%s】为主题创作一首歌曲，歌曲时间不要太短，3分钟左右，不要输出任何解释性的内容。 +输出格式如下： +歌曲名称 +第一节： +{{歌词内容}} +副歌： +{{歌词内容}} + +第二节： +{{歌词内容}} +副歌： +{{歌词内容}} + +尾声： +{{歌词内容}} +` diff --git a/api/service/video/luma.go b/api/service/video/luma.go index 144269f7..f5eba330 100644 --- a/api/service/video/luma.go +++ b/api/service/video/luma.go @@ -87,7 +87,7 @@ func (s *Service) Run() { // translate prompt if utils.HasChinese(task.Prompt) { - content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), "gpt-4o-mini", 0) + content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId) if err == nil { task.Prompt = content } else { diff --git a/api/store/model/sd_job.go b/api/store/model/sd_job.go index 8542c307..c2ff59a8 100644 --- a/api/store/model/sd_job.go +++ b/api/store/model/sd_job.go @@ -7,6 +7,7 @@ type SdJob struct { Type string UserId int TaskId string + TaskInfo string // 原始任务信息 ImgURL string Progress int Prompt string diff --git a/api/store/vo/mj_job.go b/api/store/vo/mj_job.go index 458e9a2d..ab6d99d9 100644 --- a/api/store/vo/mj_job.go +++ b/api/store/vo/mj_job.go @@ -1,21 +1,20 @@ package vo type MidJourneyJob struct { - Id uint `json:"id"` - Type string `json:"type"` - UserId int `json:"user_id"` - ChannelId string `json:"channel_id"` - TaskId string `json:"task_id"` - MessageId string `json:"message_id"` - ReferenceId string `json:"reference_id"` - ImgURL string `json:"img_url"` - OrgURL string `json:"org_url"` - Hash string `json:"hash"` - Progress int `json:"progress"` - Prompt string `json:"prompt"` - UseProxy bool `json:"use_proxy"` - Publish bool `json:"publish"` - ErrMsg string `json:"err_msg"` - Power int `json:"power"` - CreatedAt int64 `json:"created_at"` + Id uint `json:"id"` + Type string `json:"type"` + UserId int `json:"user_id"` + ChannelId string `json:"channel_id"` + TaskId string `json:"task_id"` + MessageId string `json:"message_id"` + ImgURL string `json:"img_url"` + OrgURL string `json:"org_url"` + Hash string `json:"hash"` + Progress int `json:"progress"` + Prompt string `json:"prompt"` + UseProxy bool `json:"use_proxy"` + Publish bool `json:"publish"` + ErrMsg string `json:"err_msg"` + Power int `json:"power"` + CreatedAt int64 `json:"created_at"` } diff --git a/api/utils/openai.go b/api/utils/openai.go index 3c1e4f15..06e3c465 100644 --- a/api/utils/openai.go +++ b/api/utils/openai.go @@ -45,20 +45,25 @@ type apiRes struct { } `json:"choices"` } -func OpenAIRequest(db *gorm.DB, prompt string, modelName string, keyId int) (string, error) { +func OpenAIRequest(db *gorm.DB, prompt string, modelId int) (string, error) { messages := make([]interface{}, 1) messages[0] = types.Message{ Role: "user", Content: prompt, } - return SendOpenAIMessage(db, messages, modelName, keyId) + return SendOpenAIMessage(db, messages, modelId) } -func SendOpenAIMessage(db *gorm.DB, messages []interface{}, modelName string, keyId int) (string, error) { +func SendOpenAIMessage(db *gorm.DB, messages []interface{}, modelId int) (string, error) { + var chatModel model.ChatModel + db.Where("id", modelId).First(&chatModel) + if chatModel.Name == "" { + chatModel.Name = "gpt-4o-mini" // 默认使用 gpt-4o-mini + } var apiKey model.ApiKey session := db.Session(&gorm.Session{}).Where("type", "chat").Where("enabled", true) - if keyId > 0 { - session = session.Where("id", keyId) + if chatModel.KeyId > 0 { + session = session.Where("id", chatModel.KeyId) } err := session.First(&apiKey).Error if err != nil { @@ -71,11 +76,11 @@ func SendOpenAIMessage(db *gorm.DB, messages []interface{}, modelName string, ke client.SetProxyURL(apiKey.ApiURL) } apiURL := fmt.Sprintf("%s/v1/chat/completions", apiKey.ApiURL) - logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, modelName) + logger.Infof("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, chatModel.Name) r, err := client.R().SetHeader("Body-Type", "application/json"). SetHeader("Authorization", "Bearer "+apiKey.Value). SetBody(types.ApiRequest{ - Model: modelName, + Model: chatModel.Name, Temperature: 0.9, MaxTokens: 1024, Stream: false, diff --git a/database/update-v4.1.7.sql b/database/update-v4.1.7.sql new file mode 100644 index 00000000..4233bd26 --- /dev/null +++ b/database/update-v4.1.7.sql @@ -0,0 +1 @@ +ALTER TABLE `chatgpt_sd_jobs` ADD `task_info` TEXT NOT NULL COMMENT '任务详情' AFTER `task_id`; \ No newline at end of file diff --git a/web/src/views/Suno.vue b/web/src/views/Suno.vue index 543a1ccf..97aed467 100644 --- a/web/src/views/Suno.vue +++ b/web/src/views/Suno.vue @@ -615,7 +615,7 @@ const createLyric = () => { return showMessageError("请输入歌词描述") } isGenerating.value = true - httpPost("/api/suno/lyric", {prompt: data.value.lyrics}).then(res => { + httpPost("/api/prompt/lyric", {prompt: data.value.lyrics}).then(res => { const lines = res.data.split('\n'); data.value.title = lines.shift().replace(/\*/g,"") lines.shift() diff --git a/web/src/views/admin/SysConfig.vue b/web/src/views/admin/SysConfig.vue index 8d252698..2c3b626d 100644 --- a/web/src/views/admin/SysConfig.vue +++ b/web/src/views/admin/SysConfig.vue @@ -153,14 +153,13 @@ - +