From 1f964c74e94ac5fc52a2117abd6d278bf961a356 Mon Sep 17 00:00:00 2001 From: RockYang Date: Thu, 28 Mar 2024 09:53:41 +0800 Subject: [PATCH] feat: add mj_action_power system config item --- api/core/types/config.go | 7 ++++--- api/handler/mj_handler.go | 21 +++++++++++++++++-- api/service/mj/pool.go | 6 ++---- api/service/mj/service.go | 2 +- web/src/views/ImageMj.vue | 25 +++++++++++++++++----- web/src/views/admin/SysConfig.vue | 35 +++++++++++++++++++++++++++++-- web/src/views/mobile/ImageMj.vue | 25 ++++++++++++++++------ 7 files changed, 98 insertions(+), 23 deletions(-) diff --git a/api/core/types/config.go b/api/core/types/config.go index a0a6b62b..324c3cad 100644 --- a/api/core/types/config.go +++ b/api/core/types/config.go @@ -147,9 +147,10 @@ type SystemConfig struct { VipInfoText string `json:"vip_info_text"` // 会员页面充值说明 DefaultModels []int `json:"default_models,omitempty"` // 默认开通的 AI 模型 - MjPower int `json:"mj_power,omitempty"` // MJ 绘画消耗算力 - SdPower int `json:"sd_power,omitempty"` // SD 绘画消耗算力 - DallPower int `json:"dall_power,omitempty"` // DALLE3 绘图消耗算力 + MjPower int `json:"mj_power,omitempty"` // MJ 绘画消耗算力 + MjActionPower int `json:"mj_action_power"` // MJ 操作（放大，变换）消耗算力 + SdPower int `json:"sd_power,omitempty"` // SD 绘画消耗算力 + DallPower int `json:"dall_power,omitempty"` // DALLE3 绘图消耗算力 WechatCardURL string `json:"wechat_card_url,omitempty"` // 微信客服地址 diff --git a/api/handler/mj_handler.go b/api/handler/mj_handler.go index 52526ad2..d2f9a107 100644 --- a/api/handler/mj_handler.go +++ b/api/handler/mj_handler.go @@ -245,6 +245,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) { TaskId: taskId, Progress: 0, Prompt: data.Prompt, + Power: h.App.SysConfig.MjActionPower, CreatedAt: time.Now(), } if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 { @@ -268,7 +269,23 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) { if client != nil { _ = client.Send([]byte("Task Updated")) } - + // update user's power + tx := h.DB.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power - ?", job.Power)) + // 记录算力变化日志 + if tx.Error == nil && tx.RowsAffected > 0 { + user, _ := h.GetLoginUser(c) + h.DB.Create(&model.PowerLog{ + UserId: user.Id, + Username: user.Username, + Type: types.PowerConsume, + Amount: job.Power, + Balance: user.Power - job.Power, + Mark: types.PowerSub, + Model: "mid-journey", + Remark: fmt.Sprintf("Upscale 操作，任务ID：%s", job.TaskId), + CreatedAt: time.Now(), + }) + } resp.SUCCESS(c) } @@ -295,7 +312,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) { TaskId: taskId, Progress: 0, Prompt: data.Prompt, - Power: h.App.SysConfig.MjPower, + Power: h.App.SysConfig.MjActionPower, CreatedAt: time.Now(), } if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 { diff --git a/api/service/mj/pool.go b/api/service/mj/pool.go index 2cf23ac9..14ea1da4 100644 --- a/api/service/mj/pool.go +++ b/api/service/mj/pool.go @@ -162,11 +162,9 @@ func (p *ServicePool) SyncTaskProgress() { for _, job := range items { // 失败或者 30 分钟还没完成的任务删除并退回算力 if time.Now().Sub(job.CreatedAt) > time.Minute*30 || job.Progress == -1 { + // 删除任务 p.db.Delete(&job) - // 略过 Upscale 任务 - if job.Type != types.TaskUpscale.String() { - continue - } + // 退回算力 tx := p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)) if tx.Error == nil && tx.RowsAffected > 0 { var user model.User diff --git a/api/service/mj/service.go b/api/service/mj/service.go index c739078e..861095d4 100644 --- a/api/service/mj/service.go +++ b/api/service/mj/service.go @@ -67,7 +67,7 @@ func (s *Service) Run() { continue } - // 翻译提示词 + // 如果是 mj-proxy 则自动翻译提示词 if utils.HasChinese(task.Prompt) && strings.HasPrefix(s.Name, "mj-proxy-service") { content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt)) if err == nil { diff --git a/web/src/views/ImageMj.vue b/web/src/views/ImageMj.vue index e38aa1aa..49542d48 100644 --- a/web/src/views/ImageMj.vue +++ b/web/src/views/ImageMj.vue @@ -506,7 +506,7 @@ const options = [ ] const router = useRouter() -const params = ref({ +const initParams = { task_type: "image", rate: rates[0].value, model: models[0].value, @@ -520,7 +520,8 @@ const params = ref({ neg_prompt: "", tile: false, quality: 0 -}) +} +const params = ref(initParams) const imgList = ref([]) @@ -618,6 +619,15 @@ onUnmounted(() => { clipboard.value.destroy() }) +const mjPower = ref(1) +const mjActionPower = ref(1) +httpGet("/api/config/get?key=system").then(res => { + mjPower.value = res.data["mj_power"] + mjActionPower.value = res.data["mj_action_power"] +}).catch(e => { + ElMessage.error("获取系统配置失败：" + e.message) +}) + // 获取运行中的任务 const fetchRunningJobs = () => { httpGet(`/api/mj/jobs?status=0`).then(res => { @@ -632,7 +642,11 @@ const fetchRunningJobs = () => { type: 'error', duration: 0, }) - power.value += 1 + if (jobs[i].type === 'image') { + power.value += mjPower.value + } else { + power.value += mjActionPower.value + } continue } _jobs.push(jobs[i]) @@ -749,7 +763,8 @@ const generate = () => { params.value.img_arr = imgList.value httpPost("/api/mj/image", params.value).then(() => { ElMessage.success("绘画任务推送成功，请耐心等待任务执行...") - power.value -= 1 + power.value -= mjPower.value + params.value = initParams }).catch(e => { ElMessage.error("任务推送失败：" + e.message) }) @@ -775,7 +790,7 @@ const send = (url, index, item) => { prompt: item.prompt, }).then(() => { ElMessage.success("任务推送成功，请耐心等待任务执行...") - power.value -= 1 + power.value -= mjActionPower.value }).catch(e => { ElMessage.error("任务推送失败：" + e.message) }) diff --git a/web/src/views/admin/SysConfig.vue b/web/src/views/admin/SysConfig.vue index bd0331c6..f17a4eb3 100644 --- a/web/src/views/admin/SysConfig.vue +++ b/web/src/views/admin/SysConfig.vue @@ -44,8 +44,39 @@ - - + +
+ +
+ + + + + +
+
+ + +
+ +
+ + + + + +
+
diff --git a/web/src/views/mobile/ImageMj.vue b/web/src/views/mobile/ImageMj.vue index d401a70d..0c705c58 100644 --- a/web/src/views/mobile/ImageMj.vue +++ b/web/src/views/mobile/ImageMj.vue @@ -91,7 +91,7 @@
- 可用额度:{{ imgCalls }} + 可用算力额度:{{ power }} 立即生成
@@ -253,16 +253,16 @@ const params = ref({ tile: false, quality: 0 }) -const imgCalls = ref(0) const userId = ref(0) const router = useRouter() const runningJobs = ref([]) const finishedJobs = ref([]) const socket = ref(null) +const power = ref(0) onMounted(() => { checkSession().then(user => { - imgCalls.value = user['img_calls'] + power.value = user['power'] userId.value = user.id fetchRunningJobs() @@ -278,6 +278,15 @@ onUnmounted(() => { socket.value = null }) +const mjPower = ref(1) +const mjActionPower = ref(1) +httpGet("/api/config/get?key=system").then(res => { + mjPower.value = res.data["mj_power"] + mjActionPower.value = res.data["mj_action_power"] +}).catch(e => { + ElMessage.error("获取系统配置失败：" + e.message) +}) + const heartbeatHandle = ref(null) const connect = () => { let host = process.env.VUE_APP_WS_HOST @@ -335,7 +344,11 @@ const fetchRunningJobs = (userId) => { message: `任务执行失败：${jobs[i]['err_msg']}`, type: 'danger', }) - imgCalls.value += 1 + if (jobs[i].type === 'image') { + power.value += mjPower.value + } else { + power.value += mjActionPower.value + } continue } _jobs.push(jobs[i]) @@ -437,7 +450,7 @@ const send = (url, index, item) => { prompt: item.prompt, }).then(() => { ElMessage.success("任务推送成功，请耐心等待任务执行...") - imgCalls.value -= 1 + power.value -= mjActionPower.value }).catch(e => { ElMessage.error("任务推送失败：" + e.message) }) @@ -464,7 +477,7 @@ const generate = () => { params.value.img_arr = imgList.value.map(img => img.url) httpPost("/api/mj/image", params.value).then(() => { showToast("绘画任务推送成功，请耐心等待任务执行") - imgCalls.value -= 1 + power.value -= mjPower.value }).catch(e => { showFailToast("任务推送失败：" + e.message) })