refactor: use waterflow component in mj, sd and dall image drawing page

api/service/dalle/service.go

@@ -8,6 +8,8 @@ package dalle
 // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
 import (
+	"errors"
+	"fmt"
 	"geekai/core/types"
 	logger2 "geekai/logger"
 	"geekai/service"
@@ -16,8 +18,6 @@ import (
 	"geekai/store"
 	"geekai/store/model"
 	"geekai/utils"
-	"errors"
-	"fmt"
 	"github.com/go-redis/redis/v8"
 	"time"
 
@@ -261,7 +261,7 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
 	if res.Error != nil {
 		return "", err
 	}
-	s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Failed})
+	s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Finished})
 	return imgURL, nil
 }
 
@@ -294,7 +294,7 @@ func (s *Service) CheckTaskStatus() {
 							Balance:   user.Power + job.Power,
 							Mark:      types.PowerAdd,
 							Model:     "dall-e-3",
-							Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d", job.Id),
+							Remark:    fmt.Sprintf("任务失败，退回算力。任务ID：%d", job.Id),
 							CreatedAt: time.Now(),
 						})
 					}

api/service/mj/pool.go

@@ -31,7 +31,7 @@ type ServicePool struct {
 	db              *gorm.DB
 	uploaderManager *oss.UploaderManager
 	Clients         *types.LMap[uint, *types.WsClient] // UserId => Client
-	licenseService *service.LicenseService
+	licenseService  *service.LicenseService
 }
 
 var logger = logger2.GetLogger()
@@ -56,6 +56,7 @@ func (p *ServicePool) InitServices(plusConfigs []types.MjPlusConfig, proxyConfig
 	for _, s := range p.services {
 		s.Stop()
 	}
+	p.services = make([]*Service, 0)
 
 	for k, config := range plusConfigs {
 		if config.Enabled == false {

api/service/mj/service.go

@@ -28,7 +28,7 @@ type Service struct {
 	taskQueue   *store.RedisQueue
 	notifyQueue *store.RedisQueue
 	db          *gorm.DB
-	running bool
+	running     bool
 }
 
 func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, cli Client) *Service {
@@ -38,7 +38,7 @@ func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.Red
 		taskQueue:   taskQueue,
 		notifyQueue: notifyQueue,
 		Client:      cli,
-		running: true,
+		running:     true,
 	}
 }
 
@@ -129,7 +129,6 @@ func (s *Service) Run() {
 
 func (s *Service) Stop() {
 	s.running = false
-	s.Client = nil
 }
 
 type CBReq struct {

api/service/sd/pool.go

@@ -25,8 +25,8 @@ type ServicePool struct {
 	notifyQueue *store.RedisQueue
 	db          *gorm.DB
 	Clients     *types.LMap[uint, *types.WsClient] // UserId => Client
-	uploader *oss.UploaderManager
-	levelDB  *store.LevelDB
+	uploader    *oss.UploaderManager
+	levelDB     *store.LevelDB
 }
 
 func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, levelDB *store.LevelDB) *ServicePool {
@@ -40,8 +40,8 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 		services:    services,
 		db:          db,
 		Clients:     types.NewLMap[uint, *types.WsClient](),
-		uploader: manager,
-		levelDB:  levelDB,
+		uploader:    manager,
+		levelDB:     levelDB,
 	}
 }
 
@@ -50,6 +50,7 @@ func (p *ServicePool) InitServices(configs []types.StableDiffusionConfig) {
 	for _, s := range p.services {
 		s.Stop()
 	}
+	p.services = make([]*Service, 0)
 
 	for k, config := range configs {
 		if config.Enabled == false {

api/service/sd/service.go

@@ -60,7 +60,7 @@ func (s *Service) Run() {
 			logger.Errorf("taking task with error: %v", err)
 			continue
 		}
-		logger.Infof("%s handle a new Stable-Diffusion task: %+v", s.name, task)
+
 		// translate prompt
 		if utils.HasChinese(task.Params.Prompt) {
 			content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Params.Prompt))