mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-11-09 02:33:42 +08:00
feat: migrate the chatgpt-plus-ext project code to this project
This commit is contained in:
RockYang
169
api/service/sd/client.go
Normal file
169
api/service/sd/client.go
Normal file
@@ -0,0 +1,169 @@
|
||||
package sd
|
||||
|
||||
import (
|
||||
"chatplus/core/types"
|
||||
"chatplus/utils"
|
||||
"fmt"
|
||||
"github.com/imroc/req/v3"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
httpClient *req.Client
|
||||
config *types.StableDiffusionConfig
|
||||
}
|
||||
|
||||
func NewSdClient(config *types.AppConfig) *Client {
|
||||
return &Client{
|
||||
config: &config.SdConfig,
|
||||
httpClient: req.C(),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Txt2Img(params types.SdTaskParams) error {
|
||||
var data []interface{}
|
||||
err := utils.JsonDecode(Text2ImgParamTemplate, &data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
data[ParamKeys["task_id"]] = params.TaskId
|
||||
data[ParamKeys["prompt"]] = params.Prompt
|
||||
data[ParamKeys["negative_prompt"]] = params.NegativePrompt
|
||||
data[ParamKeys["steps"]] = params.Steps
|
||||
data[ParamKeys["sampler"]] = params.Sampler
|
||||
data[ParamKeys["face_fix"]] = params.FaceFix
|
||||
data[ParamKeys["cfg_scale"]] = params.CfgScale
|
||||
data[ParamKeys["seed"]] = params.Seed
|
||||
data[ParamKeys["height"]] = params.Height
|
||||
data[ParamKeys["width"]] = params.Width
|
||||
data[ParamKeys["hd_fix"]] = params.HdFix
|
||||
data[ParamKeys["hd_redraw_rate"]] = params.HdRedrawRate
|
||||
data[ParamKeys["hd_scale"]] = params.HdScale
|
||||
data[ParamKeys["hd_scale_alg"]] = params.HdScaleAlg
|
||||
data[ParamKeys["hd_sample_num"]] = params.HdSampleNum
|
||||
task := TaskInfo{
|
||||
TaskId: params.TaskId,
|
||||
Data: data,
|
||||
EventData: nil,
|
||||
FnIndex: 494,
|
||||
SessionHash: "ycaxgzm9ah",
|
||||
}
|
||||
|
||||
go func() {
|
||||
c.runTask(task, c.httpClient)
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) runTask(taskInfo TaskInfo, client *req.Client) {
|
||||
body := map[string]any{
|
||||
"data": taskInfo.Data,
|
||||
"event_data": taskInfo.EventData,
|
||||
"fn_index": taskInfo.FnIndex,
|
||||
"session_hash": taskInfo.SessionHash,
|
||||
}
|
||||
|
||||
var result = make(chan CBReq)
|
||||
go func() {
|
||||
var res struct {
|
||||
Data []interface{} `json:"data"`
|
||||
IsGenerating bool `json:"is_generating"`
|
||||
Duration float64 `json:"duration"`
|
||||
AverageDuration float64 `json:"average_duration"`
|
||||
}
|
||||
var cbReq = CBReq{TaskId: taskInfo.TaskId}
|
||||
response, err := client.R().SetBody(body).SetSuccessResult(&res).Post(c.config.ApiURL + "/run/predict")
|
||||
if err != nil {
|
||||
cbReq.Message = "error with send request: " + err.Error()
|
||||
cbReq.Success = false
|
||||
result <- cbReq
|
||||
return
|
||||
}
|
||||
|
||||
if response.IsErrorState() {
|
||||
bytes, _ := io.ReadAll(response.Body)
|
||||
cbReq.Message = "error http status code: " + string(bytes)
|
||||
cbReq.Success = false
|
||||
result <- cbReq
|
||||
return
|
||||
}
|
||||
|
||||
var images []struct {
|
||||
Name string `json:"name"`
|
||||
Data interface{} `json:"data"`
|
||||
IsFile bool `json:"is_file"`
|
||||
}
|
||||
err = utils.ForceCovert(res.Data[0], &images)
|
||||
if err != nil {
|
||||
cbReq.Message = "error with decode image:" + err.Error()
|
||||
cbReq.Success = false
|
||||
result <- cbReq
|
||||
return
|
||||
}
|
||||
|
||||
var info map[string]any
|
||||
err = utils.JsonDecode(utils.InterfaceToString(res.Data[1]), &info)
|
||||
if err != nil {
|
||||
cbReq.Message = err.Error()
|
||||
cbReq.Success = false
|
||||
result <- cbReq
|
||||
return
|
||||
}
|
||||
|
||||
//for k, v := range info {
|
||||
// fmt.Println(k, " => ", v)
|
||||
//}
|
||||
cbReq.ImageName = images[0].Name
|
||||
cbReq.Seed = utils.InterfaceToString(info["seed"])
|
||||
cbReq.Success = true
|
||||
cbReq.Progress = 100
|
||||
result <- cbReq
|
||||
close(result)
|
||||
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case value := <-result:
|
||||
if value.Success {
|
||||
logger.Infof("%s/file=%s", c.config.ApiURL, value.ImageName)
|
||||
}
|
||||
return
|
||||
default:
|
||||
var progressReq = map[string]any{
|
||||
"id_task": taskInfo.TaskId,
|
||||
"id_live_preview": 1,
|
||||
}
|
||||
|
||||
var progressRes struct {
|
||||
Active bool `json:"active"`
|
||||
Queued bool `json:"queued"`
|
||||
Completed bool `json:"completed"`
|
||||
Progress float64 `json:"progress"`
|
||||
Eta float64 `json:"eta"`
|
||||
LivePreview string `json:"live_preview"`
|
||||
IDLivePreview int `json:"id_live_preview"`
|
||||
TextInfo interface{} `json:"textinfo"`
|
||||
}
|
||||
response, err := client.R().SetBody(progressReq).SetSuccessResult(&progressRes).Post(c.config.ApiURL + "/internal/progress")
|
||||
var cbReq = CBReq{TaskId: taskInfo.TaskId, Success: true}
|
||||
if err != nil { // TODO: 这里可以考虑设置失败重试次数
|
||||
logger.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
if response.IsErrorState() {
|
||||
bytes, _ := io.ReadAll(response.Body)
|
||||
logger.Error(string(bytes))
|
||||
return
|
||||
}
|
||||
|
||||
cbReq.ImageData = progressRes.LivePreview
|
||||
cbReq.Progress = int(progressRes.Progress * 100)
|
||||
fmt.Println("Progress: ", progressRes.Progress)
|
||||
fmt.Println("Image: ", progressRes.LivePreview)
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
}
|
||||
}
|
||||
72
api/service/sd/sd_service.go
Normal file
72
api/service/sd/sd_service.go
Normal file
@@ -0,0 +1,72 @@
|
||||
package sd
|
||||
|
||||
import (
|
||||
"chatplus/core/types"
|
||||
"chatplus/service/mj"
|
||||
"chatplus/store"
|
||||
"chatplus/store/model"
|
||||
"chatplus/utils"
|
||||
"context"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"gorm.io/gorm"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SD 绘画服务
|
||||
|
||||
const RunningJobKey = "StableDiffusion_Running_Job"
|
||||
|
||||
type Service struct {
|
||||
taskQueue *store.RedisQueue
|
||||
redis *redis.Client
|
||||
db *gorm.DB
|
||||
Client *Client
|
||||
}
|
||||
|
||||
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client) *Service {
|
||||
return &Service{
|
||||
redis: redisCli,
|
||||
db: db,
|
||||
Client: client,
|
||||
taskQueue: store.NewRedisQueue("stable_diffusion_task_queue", redisCli),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Run() {
|
||||
logger.Info("Starting StableDiffusion job consumer.")
|
||||
ctx := context.Background()
|
||||
for {
|
||||
_, err := s.redis.Get(ctx, RunningJobKey).Result()
|
||||
if err == nil { // 队列串行执行
|
||||
time.Sleep(time.Second * 3)
|
||||
continue
|
||||
}
|
||||
var task types.SdTask
|
||||
err = s.taskQueue.LPop(&task)
|
||||
if err != nil {
|
||||
logger.Errorf("taking task with error: %v", err)
|
||||
continue
|
||||
}
|
||||
logger.Infof("Consuming Task: %+v", task)
|
||||
err = s.Client.Txt2Img(task.Params)
|
||||
if err != nil {
|
||||
logger.Error("绘画任务执行失败：", err)
|
||||
if task.RetryCount <= 5 {
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
task.RetryCount += 1
|
||||
time.Sleep(time.Second * 3)
|
||||
continue
|
||||
}
|
||||
|
||||
// 更新任务的执行状态
|
||||
s.db.Model(&model.MidJourneyJob{}).Where("id = ?", task.Id).UpdateColumn("started", true)
|
||||
// 锁定任务执行通道，直到任务超时（5分钟）
|
||||
s.redis.Set(ctx, mj.RunningJobKey, utils.JsonEncode(task), time.Minute*5)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) PushTask(task types.SdTask) {
|
||||
logger.Infof("add a new MidJourney Task: %+v", task)
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
234
api/service/sd/types.go
Normal file
234
api/service/sd/types.go
Normal file
@@ -0,0 +1,234 @@
|
||||
package sd
|
||||
|
||||
import logger2 "chatplus/logger"
|
||||
|
||||
var logger = logger2.GetLogger()
|
||||
|
||||
type TaskInfo struct {
|
||||
TaskId string `json:"task_id"`
|
||||
Data interface{} `json:"data"`
|
||||
EventData interface{} `json:"event_data"`
|
||||
FnIndex int `json:"fn_index"`
|
||||
SessionHash string `json:"session_hash"`
|
||||
}
|
||||
|
||||
type CBReq struct {
|
||||
TaskId string
|
||||
ImageName string
|
||||
ImageData string
|
||||
Progress int
|
||||
Seed string
|
||||
Success bool
|
||||
Message string
|
||||
}
|
||||
|
||||
var ParamKeys = map[string]int{
|
||||
"task_id": 0,
|
||||
"prompt": 1,
|
||||
"negative_prompt": 2,
|
||||
"steps": 4,
|
||||
"sampler": 5,
|
||||
"face_fix": 6,
|
||||
"cfg_scale": 10,
|
||||
"seed": 11,
|
||||
"height": 17,
|
||||
"width": 18,
|
||||
"hd_fix": 19,
|
||||
"hd_redraw_rate": 20, //高清修复重绘幅度
|
||||
"hd_scale": 21, // 高清修复放大倍数
|
||||
"hd_scale_alg": 22, // 高清修复放大算法
|
||||
"hd_sample_num": 23, // 高清修复采样次数
|
||||
}
|
||||
|
||||
const Text2ImgParamTemplate = `[
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
[],
|
||||
30,
|
||||
"DPM++ SDE Karras",
|
||||
false,
|
||||
false,
|
||||
1,
|
||||
1,
|
||||
7.5,
|
||||
-1,
|
||||
-1,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
false,
|
||||
512,
|
||||
512,
|
||||
true,
|
||||
0.7,
|
||||
2,
|
||||
"Latent",
|
||||
10,
|
||||
0,
|
||||
0,
|
||||
"Use same sampler",
|
||||
"",
|
||||
"",
|
||||
[],
|
||||
"None",
|
||||
false,
|
||||
"MultiDiffusion",
|
||||
false,
|
||||
true,
|
||||
1024,
|
||||
1024,
|
||||
96,
|
||||
96,
|
||||
48,
|
||||
4,
|
||||
"None",
|
||||
2,
|
||||
false,
|
||||
10,
|
||||
1,
|
||||
1,
|
||||
64,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
0.4,
|
||||
0.4,
|
||||
0.2,
|
||||
0.2,
|
||||
"",
|
||||
"",
|
||||
"Background",
|
||||
0.2,
|
||||
-1,
|
||||
false,
|
||||
3072,
|
||||
192,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
null,
|
||||
null,
|
||||
null,
|
||||
false,
|
||||
"",
|
||||
0.5,
|
||||
true,
|
||||
false,
|
||||
"",
|
||||
"Lerp",
|
||||
false,
|
||||
"🔄",
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
"positive",
|
||||
"comma",
|
||||
0,
|
||||
false,
|
||||
false,
|
||||
"",
|
||||
"Seed",
|
||||
"",
|
||||
[],
|
||||
"Nothing",
|
||||
"",
|
||||
[],
|
||||
"Nothing",
|
||||
"",
|
||||
[],
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
0,
|
||||
null,
|
||||
null,
|
||||
false,
|
||||
null,
|
||||
null,
|
||||
false,
|
||||
null,
|
||||
null,
|
||||
false,
|
||||
50
|
||||
]`
|
||||
Reference in New Issue
Block a user