enable to update AI Drawing configuarations in admin console page

api/service/sd/pool.go

@@ -25,28 +25,14 @@ type ServicePool struct {
 	notifyQueue *store.RedisQueue
 	db          *gorm.DB
 	Clients     *types.LMap[uint, *types.WsClient] // UserId => Client
+	uploader *oss.UploaderManager
+	levelDB  *store.LevelDB
 }
 
-func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig, levelDB *store.LevelDB) *ServicePool {
+func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, levelDB *store.LevelDB) *ServicePool {
 	services := make([]*Service, 0)
 	taskQueue := store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli)
 	notifyQueue := store.NewRedisQueue("StableDiffusion_Queue", redisCli)
-	// create mj client and service
-	for _, config := range appConfig.SdConfigs {
-		if config.Enabled == false {
-			continue
-		}
-
-		// create sd service
-		name := fmt.Sprintf("StableDifffusion Service-%s", config.Model)
-		service := NewService(name, config, taskQueue, notifyQueue, db, manager, levelDB)
-		// run sd service
-		go func() {
-			service.Run()
-		}()
-
-		services = append(services, service)
-	}
 
 	return &ServicePool{
 		taskQueue:   taskQueue,
@@ -54,6 +40,31 @@ 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,
+	}
+}
+
+func (p *ServicePool) InitServices(configs []types.StableDiffusionConfig) {
+	// stop old service
+	for _, s := range p.services {
+		s.Stop()
+	}
+
+	for k, config := range configs {
+		if config.Enabled == false {
+			continue
+		}
+
+		// create sd service
+		name := fmt.Sprintf(" sd-service-%d", k)
+		service := NewService(name, config, p.taskQueue, p.notifyQueue, p.db, p.uploader, p.levelDB)
+		// run sd service
+		go func() {
+			service.Run()
+		}()
+
+		p.services = append(p.services, service)
 	}
 }
 

api/service/sd/service.go

@@ -33,6 +33,7 @@ type Service struct {
 	uploadManager *oss.UploaderManager
 	name          string // service name
 	leveldb       *store.LevelDB
+	running       bool // 运行状态
 }
 
 func NewService(name string, config types.StableDiffusionConfig, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelDB) *Service {
@@ -46,18 +47,20 @@ func NewService(name string, config types.StableDiffusionConfig, taskQueue *stor
 		db:            db,
 		leveldb:       levelDB,
 		uploadManager: manager,
+		running:       true,
 	}
 }
 
 func (s *Service) Run() {
-	for {
+	logger.Infof("Starting Stable-Diffusion job consumer for %s", s.name)
+	for s.running {
 		var task types.SdTask
 		err := s.taskQueue.LPop(&task)
 		if err != nil {
 			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))
@@ -94,6 +97,10 @@ func (s *Service) Run() {
 	}
 }
 
+func (s *Service) Stop() {
+	s.running = false
+}
+
 // Txt2ImgReq 文生图请求实体
 type Txt2ImgReq struct {
 	Prompt            string  `json:"prompt"`
@@ -104,6 +111,7 @@ type Txt2ImgReq struct {
 	Width             int     `json:"width"`
 	Height            int     `json:"height"`
 	SamplerName       string  `json:"sampler_name"`
+	Scheduler         string  `json:"scheduler"`
 	EnableHr          bool    `json:"enable_hr,omitempty"`
 	HrScale           int     `json:"hr_scale,omitempty"`
 	HrUpscaler        string  `json:"hr_upscaler,omitempty"`
@@ -137,6 +145,7 @@ func (s *Service) Txt2Img(task types.SdTask) error {
 		Width:          task.Params.Width,
 		Height:         task.Params.Height,
 		SamplerName:    task.Params.Sampler,
+		Scheduler:      task.Params.Scheduler,
 		ForceTaskId:    task.Params.TaskId,
 	}
 	if task.Params.Seed > 0 {