From 20943ad3d3e0c2d9b01589ac3cfd4bc2ec82b12b Mon Sep 17 00:00:00 2001 From: RockYang Date: Fri, 19 Jan 2024 17:10:52 +0800 Subject: [PATCH] fix: fixed bug for img_call increased when upscale task run failed --- api/service/mj/plus/service.go | 4 +++- api/service/mj/pool.go | 6 ++++-- api/service/mj/service.go | 4 +++- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/api/service/mj/plus/service.go b/api/service/mj/plus/service.go index 2a6a0435..f77f6271 100644 --- a/api/service/mj/plus/service.go +++ b/api/service/mj/plus/service.go @@ -86,7 +86,9 @@ func (s *Service) Run() { // 任务失败，通知前端 s.notifyQueue.RPush(task.UserId) // restore img_call quota - s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + if task.Type.String() != types.TaskUpscale.String() { + s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + } // TODO: 任务提交失败，加入队列重试 continue diff --git a/api/service/mj/pool.go b/api/service/mj/pool.go index 5c33d188..294a7646 100644 --- a/api/service/mj/pool.go +++ b/api/service/mj/pool.go @@ -197,8 +197,10 @@ func (p *ServicePool) SyncTaskProgress() { // 30 分钟还没完成的任务直接删除 if time.Now().Sub(v.CreatedAt) > time.Minute*30 { p.db.Delete(&v) - // 退回绘图次数 - p.db.Model(&model.User{}).Where("id = ?", v.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + // 非放大任务，退回绘图次数 + if v.Type != types.TaskUpscale.String() { + p.db.Model(&model.User{}).Where("id = ?", v.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + } continue } diff --git a/api/service/mj/service.go b/api/service/mj/service.go index 8ec8dcac..0b332807 100644 --- a/api/service/mj/service.go +++ b/api/service/mj/service.go @@ -81,7 +81,9 @@ func (s *Service) Run() { s.db.Model(&model.MidJourneyJob{Id: uint(task.Id)}).UpdateColumn("progress", -1) s.notifyQueue.RPush(task.UserId) // restore img_call quota - s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + if task.Type.String() != types.TaskUpscale.String() { + s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1)) + } continue }