xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-08 02:03:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

save task origin info for AI generating jobs

Browse Source
This commit is contained in:
RockYang
2024-11-11 17:22:08 +08:00
parent 95a071014c
commit c39814ce2b
31 changed files with 358 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
api/handler/dalle_handler.go
View File

@@ -72,10 +72,21 @@ func (h *DallJobHandler) Image(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
task := types.DallTask{
ClientId: data.ClientId,
UserId: uint(userId),
Prompt: data.Prompt,
Quality: data.Quality,
Size: data.Size,
Style: data.Style,
Power: h.App.SysConfig.DallPower,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
job := model.DallJob{
UserId: uint(userId),
Prompt: data.Prompt,
Power: h.App.SysConfig.DallPower,
UserId: uint(userId),
Prompt: data.Prompt,
Power: task.Power,
TaskInfo: utils.JsonEncode(task),
}
res := h.DB.Create(&job)
if res.Error != nil {
@@ -83,17 +94,8 @@ func (h *DallJobHandler) Image(c *gin.Context) {
return
}
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,
TranslateModelId: h.App.SysConfig.TranslateModelId,
})
task.Id = job.Id
h.dallService.PushTask(task)
resp.SUCCESS(c)
}

2
api/handler/function_handler.go
View File

@@ -212,7 +212,7 @@ func (h *FunctionHandler) Dall3(c *gin.Context) {
}
content, err := h.dallService.Image(types.DallTask{
JobId: job.Id,
Id: job.Id,
UserId: user.Id,
Prompt: job.Prompt,
N: 1,

76
api/handler/mj_handler.go
View File

@@ -152,10 +152,23 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
resp.ERROR(c, "error with generate task id: "+err.Error())
return
}
task := types.MjTask{
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,
}
job := model.MidJourneyJob{
Type: data.TaskType,
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Prompt: fmt.Sprintf("%s %s", data.Prompt, params),
Power: h.App.SysConfig.MjPower,
@@ -175,19 +188,8 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
return
}
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,
TranslateModelId: h.App.SysConfig.TranslateModelId,
})
task.Id = job.Id
h.mjService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -226,10 +228,21 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
taskId, _ := h.snowflake.Next(true)
task := types.MjTask{
ClientId: data.ClientId,
Type: types.TaskUpscale,
UserId: userId,
ChannelId: data.ChannelId,
Index: data.Index,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
}
job := model.MidJourneyJob{
Type: types.TaskUpscale.String(),
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
@@ -239,17 +252,8 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
return
}
h.mjService.PushTask(types.MjTask{
Id: job.Id,
ClientId: data.ClientId,
Type: types.TaskUpscale,
UserId: userId,
ChannelId: data.ChannelId,
Index: data.Index,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
task.Id = job.Id
h.mjService.PushTask(task)
// update user's power
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
@@ -280,11 +284,22 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
taskId, _ := h.snowflake.Next(true)
task := types.MjTask{
Type: types.TaskVariation,
ClientId: data.ClientId,
UserId: userId,
Index: data.Index,
ChannelId: data.ChannelId,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
}
job := model.MidJourneyJob{
Type: types.TaskVariation.String(),
ChannelId: data.ChannelId,
UserId: userId,
TaskId: taskId,
TaskInfo: utils.JsonEncode(task),
Progress: 0,
Power: h.App.SysConfig.MjActionPower,
CreatedAt: time.Now(),
@@ -294,17 +309,8 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
return
}
h.mjService.PushTask(types.MjTask{
Id: job.Id,
Type: types.TaskVariation,
ClientId: data.ClientId,
UserId: userId,
Index: data.Index,
ChannelId: data.ChannelId,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
Mode: h.App.SysConfig.MjMode,
})
task.Id = job.Id
h.mjService.PushTask(task)
err := h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,

37
api/handler/prompt_handler.go
View File

@@ -18,6 +18,7 @@ import (
"geekai/utils/resp"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"strings"
)
// 提示词生成 handler
@@ -57,3 +58,39 @@ func (h *PromptHandler) Lyric(c *gin.Context) {
resp.SUCCESS(c, content)
}
// Image 生成 AI 绘画提示词
func (h *PromptHandler) Image(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.ImagePromptOptimizeTemplate, data.Prompt), h.App.SysConfig.TranslateModelId)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, strings.Trim(content, `"`))
}
// Video 生成视频提示词
func (h *PromptHandler) Video(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.VideoPromptTemplate, data.Prompt), h.App.SysConfig.TranslateModelId)
if err != nil {
resp.ERROR(c, err.Error())
return
}
resp.SUCCESS(c, strings.Trim(content, `"`))
}

37
api/handler/suno_handler.go
View File

@@ -89,13 +89,29 @@ func (h *SunoHandler) Create(c *gin.Context) {
data.Prompt = fmt.Sprintf("%s\n%s", song.Prompt, refSong.Prompt)
}
}
task := types.SunoTask{
ClientId: data.ClientId,
UserId: int(h.GetLoginUserId(c)),
Type: data.Type,
Title: data.Title,
RefTaskId: data.RefTaskId,
RefSongId: data.RefSongId,
ExtendSecs: data.ExtendSecs,
Prompt: data.Prompt,
Tags: data.Tags,
Model: data.Model,
Instrumental: data.Instrumental,
SongId: data.SongId,
AudioURL: data.AudioURL,
}
// 插入数据库
job := model.SunoJob{
UserId: int(h.GetLoginUserId(c)),
UserId: task.UserId,
Prompt: data.Prompt,
Instrumental: data.Instrumental,
ModelName: data.Model,
TaskInfo: utils.JsonEncode(task),
Tags: data.Tags,
Title: data.Title,
Type: data.Type,
@@ -115,26 +131,13 @@ func (h *SunoHandler) Create(c *gin.Context) {
}
// 创建任务
h.sunoService.PushTask(types.SunoTask{
ClientId: data.ClientId,
Id: job.Id,
UserId: job.UserId,
Type: job.Type,
Title: job.Title,
RefTaskId: data.RefTaskId,
RefSongId: data.RefSongId,
ExtendSecs: data.ExtendSecs,
Prompt: job.Prompt,
Tags: data.Tags,
Model: data.Model,
Instrumental: data.Instrumental,
SongId: data.SongId,
AudioURL: data.AudioURL,
})
task.Id = job.Id
h.sunoService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerConsume,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 文生歌曲，%s", job.ModelName),
CreatedAt: time.Now(),
})

29
api/handler/video_handler.go
View File

@@ -80,13 +80,21 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
StartImgURL: data.FirstFrameImg,
EndImgURL: data.EndFrameImg,
}
task := types.VideoTask{
ClientId: data.ClientId,
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Params: params,
TranslateModelId: h.App.SysConfig.TranslateModelId,
}
// 插入数据库
job := model.VideoJob{
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Power: h.App.SysConfig.LumaPower,
Params: utils.JsonEncode(params),
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
Power: h.App.SysConfig.LumaPower,
TaskInfo: utils.JsonEncode(task),
}
tx := h.DB.Create(&job)
if tx.Error != nil {
@@ -95,15 +103,8 @@ 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,
TranslateModelId: h.App.SysConfig.TranslateModelId,
})
task.Id = job.Id
h.videoService.PushTask(task)
// update user's power
err = h.userService.DecreasePower(job.UserId, job.Power, model.PowerLog{