xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-08 02:03:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor stable diffusion service, use api key instead of configs

Browse Source
This commit is contained in:
RockYang
2024-08-07 17:30:59 +08:00
parent f9b809801d
commit c9d0700fd9
7 changed files with 139 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
api/handler/admin/config_handler.go
View File

@@ -12,7 +12,6 @@ import (
"geekai/core/types"
"geekai/handler"
"geekai/service"
"geekai/service/sd"
"geekai/store"
"geekai/store/model"
"geekai/utils"
@@ -27,14 +26,12 @@ type ConfigHandler struct {
handler.BaseHandler
levelDB *store.LevelDB
licenseService *service.LicenseService
sdServicePool *sd.ServicePool
}
func NewConfigHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB, licenseService *service.LicenseService, sdPool *sd.ServicePool) *ConfigHandler {
func NewConfigHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB, licenseService *service.LicenseService) *ConfigHandler {
return &ConfigHandler{
BaseHandler: handler.BaseHandler{App: app, DB: db},
levelDB: levelDB,
sdServicePool: sdPool,
licenseService: licenseService,
}
}

22
api/handler/sd_handler.go
View File

@@ -32,15 +32,15 @@ import (
type SdJobHandler struct {
BaseHandler
redis *redis.Client
pool *sd.ServicePool
service *sd.Service
uploader *oss.UploaderManager
snowflake *service.Snowflake
leveldb *store.LevelDB
}
func NewSdJobHandler(app *core.AppServer, db *gorm.DB, pool *sd.ServicePool, manager *oss.UploaderManager, snowflake *service.Snowflake, levelDB *store.LevelDB) *SdJobHandler {
func NewSdJobHandler(app *core.AppServer, db *gorm.DB, service *sd.Service, manager *oss.UploaderManager, snowflake *service.Snowflake, levelDB *store.LevelDB) *SdJobHandler {
return &SdJobHandler{
pool: pool,
service: service,
uploader: manager,
snowflake: snowflake,
leveldb: levelDB,
@@ -68,7 +68,7 @@ func (h *SdJobHandler) Client(c *gin.Context) {
}
client := types.NewWsClient(ws)
h.pool.Clients.Put(uint(userId), client)
h.service.Clients.Put(uint(userId), client)
logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
}
@@ -79,11 +79,6 @@ func (h *SdJobHandler) preCheck(c *gin.Context) bool {
return false
}
if !h.pool.HasAvailableService() {
resp.ERROR(c, "Stable-Diffusion 池子中没有没有可用的服务！")
return false
}
if user.Power < h.App.SysConfig.SdPower {
resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
return false
@@ -164,14 +159,14 @@ func (h *SdJobHandler) Image(c *gin.Context) {
return
}
h.pool.PushTask(types.SdTask{
h.service.PushTask(types.SdTask{
Id: int(job.Id),
Type: types.TaskImage,
Params: params,
UserId: userId,
})
client := h.pool.Clients.Get(uint(job.UserId))
client := h.service.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte("Task Updated"))
}
@@ -328,11 +323,6 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
logger.Error("remove image failed: ", err)
}
client := h.pool.Clients.Get(uint(job.UserId))
if client != nil {
_ = client.Send([]byte(service.TaskStatusFinished))
}
resp.SUCCESS(c)
}