xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2026-05-10 11:44:28 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

add user lock for chat api, Prevent insufficient deduction of user power caused by submitting multiple requests at one time

Browse Source
This commit is contained in:
GeekMaster
2025-09-12 15:05:14 +08:00
parent 65fb58585c
commit c5badb3e13
43 changed files with 309 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
api/core/types/jimeng.go
View File

@@ -14,3 +14,50 @@ type JimengPower struct {
VirtualHuman int `json:"virtual_human"` // 数字人视频生成算力，单位：积分/秒
ActionTransfer int `json:"action_transfer"` // 视频动作迁移算力，单位：积分/秒
}
// JMTaskStatus 任务状态
type JMTaskStatus string
const (
JMTaskStatusInQueue = JMTaskStatus("in_queue") // 任务已提交
JMTaskStatusGenerating = JMTaskStatus("generating") // 任务处理中
JMTaskStatusDone = JMTaskStatus("done") // 处理完成
JMTaskStatusNotFound = JMTaskStatus("not_found") // 任务未找到
JMTaskStatusSuccess = JMTaskStatus("success") // 任务成功
JMTaskStatusFailed = JMTaskStatus("failed") // 任务失败
JMTaskStatusExpired = JMTaskStatus("expired") // 任务过期
)
// JMTaskType 任务类型
type JMTaskType string
const (
JMTaskTypeImage = JMTaskType("image") // 文生图
JMTaskTypeVideo = JMTaskType("video") // 图生图
JMTaskTypeVirtualHuman = JMTaskType("virtual_human") // 图像编辑
JMTaskTypeActionTransfer = JMTaskType("action_transfer") // 图像特效
)
// JimengTaskRequest 即梦AI任务请求
type JimengTaskRequest struct {
ReqKey string `json:"req_key"` // 请求Key
// 公共参数
Prompt string `json:"prompt,omitempty"`
ImageUrls []string `json:"image_urls,omitempty"`
// 图片生成参数
Size string `json:"size,omitempty"`
UsePreLLM bool `json:"use_pre_llm,omitempty"`
// 视频生成参数
Duration string `json:"duration,omitempty"` // 视频时长
TemplateId string `json:"template_id,omitempty"` // 运镜模板ID
AspectRatio string `json:"aspect_ratio,omitempty"`
CameraStrength string `json:"camera_strength,omitempty"` // 运镜强度
// 数字人视频生成参数
AudioURL string `json:"audio_url,omitempty"` // 音频URL
// 视频动作迁移参数
VideoURL string `json:"video_url,omitempty"` // 动作视频URL
}

45
api/core/types/user_lock.go Normal file
View File

@@ -0,0 +1,45 @@
package types
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 "sync"
// UserLockManager 提供基于用户ID的TryLock功能，确保同一用户并发请求串行化
type UserLockManager struct {
mu sync.Mutex
locks map[uint]bool
}
func NewUserLockManager() *UserLockManager {
return &UserLockManager{mu: sync.Mutex{}, locks: make(map[uint]bool)}
}
// TryLock 尝试为指定用户加锁。若已被占用返回 false
func (m *UserLockManager) TryLock(userId uint) bool {
if userId == 0 {
return true
}
m.mu.Lock()
defer m.mu.Unlock()
if m.locks[userId] {
return false
}
m.locks[userId] = true
return true
}
// Unlock 释放指定用户的锁
func (m *UserLockManager) Unlock(userId uint) {
if userId == 0 {
return
}
m.mu.Lock()
delete(m.locks, userId)
m.mu.Unlock()
}

12
api/handler/admin/jimeng_handler.go
View File

@@ -131,7 +131,7 @@ func (h *AdminJimengHandler) BatchRemove(c *gin.Context) {
continue // 跳过不存在的
}
tx := h.DB.Begin()
if job.Status != model.JMTaskStatusSuccess && job.Power > 0 {
if job.Status != types.JMTaskStatusSuccess && job.Power > 0 {
remark := fmt.Sprintf("任务未成功，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg)
err = h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
@@ -172,7 +172,7 @@ func (h *AdminJimengHandler) BatchRemove(c *gin.Context) {
// Stats 获取统计信息
func (h *AdminJimengHandler) Stats(c *gin.Context) {
type StatResult struct {
Status model.JMTaskStatus `json:"status"`
Status types.JMTaskStatus `json:"status"`
Count int64 `json:"count"`
}
@@ -198,13 +198,13 @@ func (h *AdminJimengHandler) Stats(c *gin.Context) {
for _, stat := range stats {
result["totalTasks"] = result["totalTasks"].(int64) + stat.Count
switch stat.Status {
case model.JMTaskStatusInQueue:
case types.JMTaskStatusInQueue:
result["pendingTasks"] = stat.Count
case model.JMTaskStatusSuccess:
case types.JMTaskStatusSuccess:
result["completedTasks"] = stat.Count
case model.JMTaskStatusGenerating:
case types.JMTaskStatusGenerating:
result["processingTasks"] = stat.Count
case model.JMTaskStatusFailed:
case types.JMTaskStatusFailed:
result["failedTasks"] = stat.Count
}
}

10
api/handler/chat_handler.go
View File

@@ -69,6 +69,7 @@ type ChatHandler struct {
ReqCancelFunc *types.LMap[string, context.CancelFunc] // HttpClient 请求取消 handle function
userService *service.UserService
moderationManager *moderation.ServiceManager
userLocks *types.UserLockManager
}
func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService, userService *service.UserService, moderationManager *moderation.ServiceManager) *ChatHandler {
@@ -80,6 +81,7 @@ func NewChatHandler(app *core.AppServer, db *gorm.DB, redis *redis.Client, manag
ReqCancelFunc: types.NewLMap[string, context.CancelFunc](),
userService: userService,
moderationManager: moderationManager,
userLocks: types.NewUserLockManager(),
}
}
@@ -120,6 +122,14 @@ func (h *ChatHandler) Chat(c *gin.Context) {
return
}
// 用户级并发锁，确保同一用户同时只有一个对话请求
if !h.userLocks.TryLock(input.UserId) {
pushMessage(c, ChatEventError, "您有一个对话请求正在进行中，请稍后再试或先停止当前生成！")
c.Abort()
return
}
defer h.userLocks.Unlock(input.UserId)
ctx, cancel := context.WithCancel(c.Request.Context())
defer cancel()

221
api/handler/jimeng_handler.go
View File

@@ -50,37 +50,16 @@ func (h *JimengHandler) RegisterRoutes() {
}
}
// JimengTaskRequest 统一任务请求结构体
// 支持所有生图和生成视频类型
type JimengTaskRequest struct {
TaskType string `json:"task_type" binding:"required"`
Prompt string `json:"prompt"`
ImageInput string `json:"image_input"`
ImageUrls []string `json:"image_urls"`
BinaryDataBase64 []string `json:"binary_data_base64"`
Scale float64 `json:"scale"`
Width int `json:"width"`
Height int `json:"height"`
Gpen float64 `json:"gpen"`
Skin float64 `json:"skin"`
SkinUnifi float64 `json:"skin_unifi"`
GenMode string `json:"gen_mode"`
Seed int64 `json:"seed"`
UsePreLLM bool `json:"use_pre_llm"`
TemplateId string `json:"template_id"`
AspectRatio string `json:"aspect_ratio"`
}
// CreateTask 统一任务创建接口
func (h *JimengHandler) CreateTask(c *gin.Context) {
var req JimengTaskRequest
var req types.JimengTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
resp.ERROR(c, types.InvalidArgs)
return
}
// 文本审核
if h.App.SysConfig.Moderation.Enable {
if h.App.SysConfig.Moderation.Enable && req.Prompt != "" {
moderationResult, err := h.moderationManager.GetService().Moderate(req.Prompt)
if err != nil {
logger.Error("failed to moderate content: ", err)
@@ -103,166 +82,46 @@ func (h *JimengHandler) CreateTask(c *gin.Context) {
}
// 新增：除图像特效外，其他任务类型必须有提示词
if req.TaskType != "image_effects" && req.Prompt == "" {
resp.ERROR(c, "提示词不能为空")
if req.Prompt == "" && len(req.ImageUrls) == 0 {
resp.ERROR(c, "提示词和图片不能同时为空")
return
}
user, err := h.GetLoginUser(c)
if err != nil {
resp.NotAuth(c)
return
}
if req.Width == 0 {
req.Width = 1328
}
if req.Height == 0 {
req.Height = 1328
}
if req.Seed == 0 {
req.Seed = -1
}
// 获取算力消耗
var powerCost int
var taskType model.JMTaskType
var params map[string]any
var reqKey string
var modelName string
// if user.Power < powerCost {
// resp.ERROR(c, fmt.Sprintf("算力不足，需要%d算力", powerCost))
// return
// }
switch req.TaskType {
case "text_to_image":
powerCost = h.getPowerFromConfig(model.JMTaskTypeImage)
taskType = model.JMTaskTypeImage
reqKey = jimeng.ReqKeyTextToImage
modelName = "即梦文生图"
if req.Scale == 0 {
req.Scale = 2.5
}
params = map[string]any{
"seed": req.Seed,
"scale": req.Scale,
"width": req.Width,
"height": req.Height,
"use_pre_llm": req.UsePreLLM,
}
case "image_to_image":
powerCost = h.getPowerFromConfig(model.JMTaskTypeVideo)
taskType = model.JMTaskTypeVideo
reqKey = jimeng.ReqKeyImageToImagePortrait
modelName = "即梦图生图"
if req.Gpen == 0 {
req.Gpen = 0.4
}
if req.Skin == 0 {
req.Skin = 0.3
}
if req.GenMode == "" {
if req.Prompt != "" {
req.GenMode = jimeng.GenModeCreative
} else {
req.GenMode = jimeng.GenModeReference
}
}
params = map[string]any{
"image_input": req.ImageInput,
"width": req.Width,
"height": req.Height,
"gpen": req.Gpen,
"skin": req.Skin,
"skin_unifi": req.SkinUnifi,
"gen_mode": req.GenMode,
"seed": req.Seed,
}
case "image_edit":
powerCost = h.getPowerFromConfig(model.JMTaskTypeVirtualHuman)
taskType = model.JMTaskTypeVirtualHuman
reqKey = jimeng.ReqKeyImageEdit
modelName = "即梦图像编辑"
if req.Scale == 0 {
req.Scale = 0.5
}
params = map[string]any{
"seed": req.Seed,
"scale": req.Scale,
}
params["image_urls"] = []string{req.ImageInput}
case "image_effects":
powerCost = h.getPowerFromConfig(model.JMTaskTypeActionTransfer)
taskType = model.JMTaskTypeActionTransfer
reqKey = jimeng.ReqKeyImageEffects
modelName = "即梦图像特效"
if req.Width == 0 {
req.Width = 1328
}
if req.Height == 0 {
req.Height = 1328
}
params = map[string]any{
"image_input1": req.ImageInput,
"template_id": req.TemplateId,
"width": req.Width,
"height": req.Height,
}
case "text_to_video":
powerCost = h.getPowerFromConfig(model.JMTaskTypeVideo)
taskType = model.JMTaskTypeVideo
reqKey = jimeng.ReqKeyTextToVideo
modelName = "即梦文生视频"
if req.AspectRatio == "" {
req.AspectRatio = jimeng.AspectRatio16_9
}
params = map[string]any{
"seed": req.Seed,
"aspect_ratio": req.AspectRatio,
}
case "image_to_video":
powerCost = h.getPowerFromConfig(model.JMTaskTypeVideo)
taskType = model.JMTaskTypeVideo
reqKey = jimeng.ReqKeyImageToVideo
modelName = "即梦图生视频"
params = map[string]any{
"seed": req.Seed,
"aspect_ratio": req.AspectRatio,
}
if len(req.ImageUrls) > 0 {
params["image_urls"] = req.ImageUrls
}
if len(req.BinaryDataBase64) > 0 {
params["binary_data_base64"] = req.BinaryDataBase64
}
default:
resp.ERROR(c, "不支持的任务类型")
return
}
// taskReq := &jimeng.CreateTaskRequest{
// Type: taskType,
// Prompt: req.Prompt,
// Params: params,
// ReqKey: reqKey,
// Power: powerCost,
// }
if user.Power < powerCost {
resp.ERROR(c, fmt.Sprintf("算力不足，需要%d算力", powerCost))
return
}
// job, err := h.jimengService.CreateTask(user.Id, taskReq)
// if err != nil {
// logger.Errorf("create jimeng task failed: %v", err)
// resp.ERROR(c, "创建任务失败")
// return
// }
taskReq := &jimeng.CreateTaskRequest{
Type: taskType,
Prompt: req.Prompt,
Params: params,
ReqKey: reqKey,
Power: powerCost,
}
// h.userService.DecreasePower(user.Id, powerCost, model.PowerLog{
// Type: types.PowerConsume,
// Model: "jimeng",
// Remark: fmt.Sprintf("%s，任务ID：%d", modelName, job.Id),
// })
job, err := h.jimengService.CreateTask(user.Id, taskReq)
if err != nil {
logger.Errorf("create jimeng task failed: %v", err)
resp.ERROR(c, "创建任务失败")
return
}
h.userService.DecreasePower(user.Id, powerCost, model.PowerLog{
Type: types.PowerConsume,
Model: "jimeng",
Remark: fmt.Sprintf("%s，任务ID：%d", modelName, job.Id),
})
resp.SUCCESS(c, job)
resp.SUCCESS(c)
}
// Jobs 获取任务列表
@@ -287,9 +146,9 @@ func (h *JimengHandler) Jobs(c *gin.Context) {
switch req.Filter {
case "image":
query = query.Where("type = ?", model.JMTaskTypeImage)
query = query.Where("type = ?", types.JMTaskTypeImage)
case "video":
query = query.Where("type = ?", model.JMTaskTypeVideo)
query = query.Where("type = ?", types.JMTaskTypeVideo)
}
if len(req.Ids) > 0 {
@@ -349,7 +208,7 @@ func (h *JimengHandler) Remove(c *gin.Context) {
}
// 正在运行中的任务不能删除
if job.Status == model.JMTaskStatusGenerating || job.Status == model.JMTaskStatusInQueue {
if job.Status == types.JMTaskStatusGenerating || job.Status == types.JMTaskStatusInQueue {
resp.ERROR(c, "正在运行中的任务不能删除，否则无法退回算力")
return
}
@@ -362,7 +221,7 @@ func (h *JimengHandler) Remove(c *gin.Context) {
}
// 失败任务删除后退回算力
if job.Status != model.JMTaskStatusFailed {
if job.Status != types.JMTaskStatusFailed {
err = h.userService.IncreasePower(user.Id, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "jimeng",
@@ -403,13 +262,13 @@ func (h *JimengHandler) Retry(c *gin.Context) {
}
// 只有失败的任务才能重试
if job.Status != model.JMTaskStatusFailed {
if job.Status != types.JMTaskStatusFailed {
resp.ERROR(c, "只有失败的任务才能重试")
return
}
// 重置任务状态
if err := h.jimengService.UpdateJobStatus(uint(jobId), model.JMTaskStatusInQueue, ""); err != nil {
if err := h.jimengService.UpdateJobStatus(uint(jobId), types.JMTaskStatusInQueue, ""); err != nil {
logger.Errorf("reset job status failed: %v", err)
resp.ERROR(c, "重置任务状态失败")
return
@@ -426,17 +285,17 @@ func (h *JimengHandler) Retry(c *gin.Context) {
}
// getPowerFromConfig 从配置中获取指定类型的算力消耗
func (h *JimengHandler) getPowerFromConfig(taskType model.JMTaskType) int {
func (h *JimengHandler) getPowerFromConfig(taskType types.JMTaskType) int {
config := h.App.SysConfig.Jimeng
switch taskType {
case model.JMTaskTypeImage:
case types.JMTaskTypeImage:
return config.Power.Image
case model.JMTaskTypeVideo:
case types.JMTaskTypeVideo:
return config.Power.Video
case model.JMTaskTypeVirtualHuman:
case types.JMTaskTypeVirtualHuman:
return config.Power.VirtualHuman
case model.JMTaskTypeActionTransfer:
case types.JMTaskTypeActionTransfer:
return config.Power.ActionTransfer
default:
return 10

39
api/service/jimeng/service.go
View File

@@ -9,6 +9,7 @@ import (
"gorm.io/gorm"
"geekai/core/types"
logger2 "geekai/logger"
"geekai/service/oss"
"geekai/store"
@@ -95,7 +96,7 @@ func (s *Service) processNextTask() {
if err := s.ProcessTask(jobId); err != nil {
logger.Errorf("process jimeng task failed: job_id=%d, error=%v", jobId, err)
s.UpdateJobStatus(jobId, model.JMTaskStatusFailed, err.Error())
s.UpdateJobStatus(jobId, types.JMTaskStatusFailed, err.Error())
} else {
logger.Infof("Jimeng task processed successfully: job_id=%d", jobId)
}
@@ -120,7 +121,7 @@ func (s *Service) CreateTask(userId uint, req *CreateTaskRequest) (*model.Jimeng
ReqKey: req.ReqKey,
Prompt: req.Prompt,
Params: string(paramsJson),
Status: model.JMTaskStatusInQueue,
Status: types.JMTaskStatusInQueue,
Power: req.Power,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
@@ -148,7 +149,7 @@ func (s *Service) ProcessTask(jobId uint) error {
}
// 更新任务状态为处理中
if err := s.UpdateJobStatus(job.Id, model.JMTaskStatusGenerating, ""); err != nil {
if err := s.UpdateJobStatus(job.Id, types.JMTaskStatusGenerating, ""); err != nil {
return fmt.Errorf("update job status failed: %w", err)
}
@@ -199,13 +200,13 @@ func (s *Service) buildTaskRequest(job *model.JimengJob) (*SubmitTaskRequest, er
// 根据任务类型设置特定参数
switch job.Type {
case model.JMTaskTypeImage:
case types.JMTaskTypeImage:
s.setTextToImageParams(req, params)
case model.JMTaskTypeVideo:
case types.JMTaskTypeVideo:
s.setImageToImageParams(req, params)
case model.JMTaskTypeVirtualHuman:
case types.JMTaskTypeVirtualHuman:
s.setImageEditParams(req, params)
case model.JMTaskTypeActionTransfer:
case types.JMTaskTypeActionTransfer:
s.setImageEffectsParams(req, params)
default:
return nil, fmt.Errorf("unsupported task type: %s", job.Type)
@@ -353,7 +354,7 @@ func (s *Service) pollTaskStatus() {
for {
var jobs []model.JimengJob
s.db.Where("status IN (?)", []model.JMTaskStatus{model.JMTaskStatusGenerating, model.JMTaskStatusInQueue}).Find(&jobs)
s.db.Where("status IN (?)", []types.JMTaskStatus{types.JMTaskStatusGenerating, types.JMTaskStatusInQueue}).Find(&jobs)
if len(jobs) == 0 {
logger.Debugf("no jimeng task to poll, sleep 10s")
time.Sleep(10 * time.Second)
@@ -389,7 +390,7 @@ func (s *Service) pollTaskStatus() {
}
switch resp.Data.Status {
case model.JMTaskStatusDone:
case types.JMTaskStatusDone:
// 判断任务是否成功
if resp.Message != "Success" {
s.handleTaskError(job.Id, fmt.Sprintf("task failed: %s", resp.Data.AlgorithmBaseResp.StatusMessage))
@@ -398,7 +399,7 @@ func (s *Service) pollTaskStatus() {
// 任务完成，更新结果
updates := map[string]any{
"status": model.JMTaskStatusSuccess,
"status": types.JMTaskStatusSuccess,
"updated_at": time.Now(),
}
@@ -421,15 +422,15 @@ func (s *Service) pollTaskStatus() {
}
s.db.Model(&model.JimengJob{}).Where("id = ?", job.Id).Updates(updates)
case model.JMTaskStatusInQueue, model.JMTaskStatusGenerating:
case types.JMTaskStatusInQueue, types.JMTaskStatusGenerating:
// 任务处理中
s.UpdateJobStatus(job.Id, model.JMTaskStatusGenerating, "")
s.UpdateJobStatus(job.Id, types.JMTaskStatusGenerating, "")
case model.JMTaskStatusNotFound:
case types.JMTaskStatusNotFound:
// 任务未找到
s.handleTaskError(job.Id, "task not found")
case model.JMTaskStatusExpired:
case types.JMTaskStatusExpired:
continue
default:
logger.Warnf("unknown task status: %s", resp.Data.Status)
@@ -444,7 +445,7 @@ func (s *Service) pollTaskStatus() {
}
// UpdateJobStatus 更新任务状态
func (s *Service) UpdateJobStatus(jobId uint, status model.JMTaskStatus, errMsg string) error {
func (s *Service) UpdateJobStatus(jobId uint, status types.JMTaskStatus, errMsg string) error {
updates := map[string]any{
"status": status,
"updated_at": time.Now(),
@@ -458,7 +459,7 @@ func (s *Service) UpdateJobStatus(jobId uint, status model.JMTaskStatus, errMsg
// handleTaskError 处理任务错误
func (s *Service) handleTaskError(jobId uint, errMsg string) error {
logger.Errorf("Jimeng task error (job_id: %d): %s", jobId, errMsg)
return s.UpdateJobStatus(jobId, model.JMTaskStatusFailed, errMsg)
return s.UpdateJobStatus(jobId, types.JMTaskStatusFailed, errMsg)
}
// PushTaskToQueue 推送任务到队列（用于手动重试）
@@ -469,8 +470,8 @@ func (s *Service) PushTaskToQueue(jobId uint) error {
// GetTaskStats 获取任务统计信息
func (s *Service) GetTaskStats() (map[string]any, error) {
type StatResult struct {
Status string `json:"status"`
Count int64 `json:"count"`
Status types.JMTaskStatus `json:"status"`
Count int64 `json:"count"`
}
var stats []StatResult
@@ -492,7 +493,7 @@ func (s *Service) GetTaskStats() (map[string]any, error) {
for _, stat := range stats {
result["total"] = result["total"].(int64) + stat.Count
result[stat.Status] = stat.Count
result[string(stat.Status)] = stat.Count
}
return result, nil

67
api/service/jimeng/types.go
View File

@@ -1,15 +1,7 @@
package jimeng
import "geekai/store/model"
// ReqKey 常量定义
const (
ReqKeyTextToImage = "high_aes_general_v30l_zt2i" // 文生图
ReqKeyImageToImagePortrait = "i2i_portrait_photo" // 图生图人像写真
ReqKeyImageEdit = "seededit_v3.0" // 图像编辑
ReqKeyImageEffects = "i2i_multi_style_zx2x" // 图像特效
ReqKeyTextToVideo = "jimeng_vgfm_t2v_l20" // 文生视频
ReqKeyImageToVideo = "jimeng_vgfm_i2v_l20" // 图生视频
import (
"geekai/core/types"
)
// SubmitTaskRequest 提交任务请求
@@ -73,7 +65,7 @@ type QueryTaskResponse struct {
ImageUrls []string `json:"image_urls"`
VideoUrl string `json:"video_url"`
RespData string `json:"resp_data"`
Status model.JMTaskStatus `json:"status"`
Status types.JMTaskStatus `json:"status"`
LlmResult string `json:"llm_result"`
PeResult string `json:"pe_result"`
PredictTagsResult string `json:"predict_tags_result"`
@@ -85,61 +77,10 @@ type QueryTaskResponse struct {
// CreateTaskRequest 创建任务请求
type CreateTaskRequest struct {
Type model.JMTaskType `json:"type"`
Type types.JMTaskType `json:"type"`
Prompt string `json:"prompt"`
Params map[string]any `json:"params"`
ReqKey string `json:"req_key"`
ImageUrls []string `json:"image_urls,omitempty"`
Power int `json:"power,omitempty"`
}
// LogoInfo 水印信息
type LogoInfo struct {
AddLogo bool `json:"add_logo"`
Position int `json:"position"`
Language int `json:"language"`
Opacity float64 `json:"opacity"`
LogoTextContent string `json:"logo_text_content"`
}
// ReqJsonConfig 查询配置
type ReqJsonConfig struct {
ReturnUrl bool `json:"return_url"`
LogoInfo *LogoInfo `json:"logo_info,omitempty"`
}
// ImageEffectTemplate 图像特效模板
const (
TemplateIdFelt3DPolaroid = "felt_3d_polaroid" // 毛毡3d拍立得风格
TemplateIdMyWorld = "my_world" // 像素世界风
TemplateIdMyWorldUniversal = "my_world_universal" // 像素世界-万物通用版
TemplateIdPlasticBubbleFigure = "plastic_bubble_figure" // 盲盒玩偶风
TemplateIdPlasticBubbleFigureCartoon = "plastic_bubble_figure_cartoon_text" // 塑料泡罩人偶-文字卡头版
TemplateIdFurryDreamDoll = "furry_dream_doll" // 毛绒玩偶风
TemplateIdMicroLandscapeMiniWorld = "micro_landscape_mini_world" // 迷你世界玩偶风
TemplateIdMicroLandscapeProfessional = "micro_landscape_mini_world_professional" // 微型景观小世界-职业版
TemplateIdAcrylicOrnaments = "acrylic_ornaments" // 亚克力挂饰
TemplateIdFeltKeychain = "felt_keychain" // 毛毡钥匙扣
TemplateIdLofiPixelCharacter = "lofi_pixel_character_mini_card" // Lofi像素人物小卡
TemplateIdAngelFigurine = "angel_figurine" // 天使形象手办
TemplateIdLyingInFluffyBelly = "lying_in_fluffy_belly" // 躺在毛茸茸肚皮里
TemplateIdGlassBall = "glass_ball" // 玻璃球
)
// AspectRatio 视频宽高比
const (
AspectRatio16_9 = "16:9" // 1280×720
AspectRatio9_16 = "9:16" // 720×1280
AspectRatio1_1 = "1:1" // 960×960
AspectRatio4_3 = "4:3" // 960×720
AspectRatio3_4 = "3:4" // 720×960
AspectRatio21_9 = "21:9" // 1680×720
AspectRatio9_21 = "9:21" // 720×1680
)
// GenMode 生成模式
const (
GenModeCreative = "creative" // 提示词模式
GenModeReference = "reference" // 全参考模式
GenModeReferenceChar = "reference_char" // 人物参考模式
)

56
api/store/model/jimeng_job.go
View File

@@ -1,52 +1,30 @@
package model
import (
"geekai/core/types"
"time"
)
// JimengJob 即梦AI任务模型
type JimengJob struct {
Id uint `gorm:"column:id;primaryKey;autoIncrement" json:"id"`
UserId uint `gorm:"column:user_id;type:int(11);not null;index;comment:用户ID" json:"user_id"`
TaskId string `gorm:"column:task_id;type:varchar(100);not null;index;comment:任务ID" json:"task_id"`
Type JMTaskType `gorm:"column:type;type:varchar(50);not null;comment:任务类型" json:"type"`
ReqKey string `gorm:"column:req_key;type:varchar(100);comment:请求Key" json:"req_key"`
Prompt string `gorm:"column:prompt;type:text;comment:提示词" json:"prompt"`
Params string `gorm:"column:params;type:text;comment:任务参数JSON" json:"params"`
ImgURL string `gorm:"column:img_url;type:varchar(1024);comment:图片或封面URL" json:"img_url"`
VideoURL string `gorm:"column:video_url;type:varchar(1024);comment:视频URL" json:"video_url"`
RawData string `gorm:"column:raw_data;type:text;comment:原始API响应" json:"raw_data"`
Progress int `gorm:"column:progress;type:int;default:0;comment:进度百分比" json:"progress"`
Status JMTaskStatus `gorm:"column:status;type:varchar(20);default:'pending';comment:任务状态" json:"status"`
ErrMsg string `gorm:"column:err_msg;type:varchar(1024);comment:错误信息" json:"err_msg"`
Power int `gorm:"column:power;type:int(11);default:0;comment:消耗算力" json:"power"`
CreatedAt time.Time `gorm:"column:created_at;type:datetime;not null;comment:创建时间" json:"created_at"`
UpdatedAt time.Time `gorm:"column:updated_at;type:datetime;not null;comment:更新时间" json:"updated_at"`
Id uint `gorm:"column:id;primaryKey;autoIncrement" json:"id"`
UserId uint `gorm:"column:user_id;type:int(11);not null;index;comment:用户ID" json:"user_id"`
TaskId string `gorm:"column:task_id;type:varchar(100);not null;index;comment:任务ID" json:"task_id"`
Type types.JMTaskType `gorm:"column:type;type:varchar(50);not null;comment:任务类型" json:"type"`
ReqKey string `gorm:"column:req_key;type:varchar(100);comment:请求Key" json:"req_key"`
Prompt string `gorm:"column:prompt;type:text;comment:提示词" json:"prompt"`
Params string `gorm:"column:params;type:text;comment:任务参数JSON" json:"params"`
ImgURL string `gorm:"column:img_url;type:varchar(1024);comment:图片或封面URL" json:"img_url"`
VideoURL string `gorm:"column:video_url;type:varchar(1024);comment:视频URL" json:"video_url"`
RawData string `gorm:"column:raw_data;type:text;comment:原始API响应" json:"raw_data"`
Progress int `gorm:"column:progress;type:int;default:0;comment:进度百分比" json:"progress"`
Status types.JMTaskStatus `gorm:"column:status;type:varchar(20);default:'pending';comment:任务状态" json:"status"`
ErrMsg string `gorm:"column:err_msg;type:varchar(1024);comment:错误信息" json:"err_msg"`
Power int `gorm:"column:power;type:int(11);default:0;comment:消耗算力" json:"power"`
CreatedAt time.Time `gorm:"column:created_at;type:datetime;not null;comment:创建时间" json:"created_at"`
UpdatedAt time.Time `gorm:"column:updated_at;type:datetime;not null;comment:更新时间" json:"updated_at"`
}
// JMTaskStatus 任务状态
type JMTaskStatus string
const (
JMTaskStatusInQueue = JMTaskStatus("in_queue") // 任务已提交
JMTaskStatusGenerating = JMTaskStatus("generating") // 任务处理中
JMTaskStatusDone = JMTaskStatus("done") // 处理完成
JMTaskStatusNotFound = JMTaskStatus("not_found") // 任务未找到
JMTaskStatusSuccess = JMTaskStatus("success") // 任务成功
JMTaskStatusFailed = JMTaskStatus("failed") // 任务失败
JMTaskStatusExpired = JMTaskStatus("expired") // 任务过期
)
// JMTaskType 任务类型
type JMTaskType string
const (
JMTaskTypeImage = JMTaskType("image") // 文生图
JMTaskTypeVideo = JMTaskType("video") // 图生图
JMTaskTypeVirtualHuman = JMTaskType("virtual_human") // 图像编辑
JMTaskTypeActionTransfer = JMTaskType("action_transfer") // 图像特效
)
// TableName 返回数据表名称
func (JimengJob) TableName() string {
return "geekai_jimeng_jobs"

6
api/store/vo/jimeng_job.go
View File

@@ -1,13 +1,13 @@
package vo
import "geekai/store/model"
import "geekai/core/types"
// JimengJob 即梦AI任务VO
type JimengJob struct {
Id uint `json:"id"`
UserId uint `json:"user_id"`
TaskId string `json:"task_id"`
Type model.JMTaskType `json:"type"`
Type types.JMTaskType `json:"type"`
ReqKey string `json:"req_key"`
Prompt string `json:"prompt"`
Params map[string]any `json:"params"`
@@ -15,7 +15,7 @@ type JimengJob struct {
VideoURL string `json:"video_url"`
RawData string `json:"raw_data"`
Progress int `json:"progress"`
Status model.JMTaskStatus `json:"status"`
Status types.JMTaskStatus `json:"status"`
ErrMsg string `json:"err_msg"`
Power int `json:"power"`
CreatedAt int64 `json:"created_at"` // 时间戳