feat: stable diffusion page is ready

2025-11-07 09:43:43 +08:00 · 2023-09-28 18:09:45 +08:00
parent f873d6b375
commit cc7271aa73
23 changed files with 1730 additions and 779 deletions
--- a/api/handler/sd_handler.go
+++ b/api/handler/sd_handler.go
@@ -1,316 +1,202 @@
 package handler

-//
-//import (
-//	"chatplus/core"
-//	"chatplus/core/types"
-//	"chatplus/service"
-//	"chatplus/service/oss"
-//	"chatplus/store/model"
-//	"chatplus/store/vo"
-//	"chatplus/utils"
-//	"chatplus/utils/resp"
-//	"encoding/base64"
-//	"fmt"
-//	"github.com/gin-gonic/gin"
-//	"github.com/go-redis/redis/v8"
-//	"github.com/gorilla/websocket"
-//	"gorm.io/gorm"
-//	"net/http"
-//	"strings"
-//	"sync"
-//	"time"
-//)
-//
-//type SdJobHandler struct {
-//	BaseHandler
-//	redis           *redis.Client
-//	db              *gorm.DB
-//	mjService       *service.MjService
-//	uploaderManager *oss.UploaderManager
-//	lock            sync.Mutex
-//	clients         *types.LMap[string, *types.WsClient]
-//}
-//
-//func NewSdJobHandler(
-//	app *core.AppServer,
-//	client *redis.Client,
-//	db *gorm.DB,
-//	manager *oss.UploaderManager,
-//	mjService *service.MjService) *MidJourneyHandler {
-//	h := MidJourneyHandler{
-//		redis:           client,
-//		db:              db,
-//		uploaderManager: manager,
-//		lock:            sync.Mutex{},
-//		mjService:       mjService,
-//		clients:         types.NewLMap[string, *types.WsClient](),
-//	}
-//	h.App = app
-//	return &h
-//}
-//
-//// Client WebSocket 客户端，用于通知任务状态变更
-//func (h *SdJobHandler) Client(c *gin.Context) {
-//	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
-//	if err != nil {
-//		logger.Error(err)
-//		return
-//	}
-//
-//	sessionId := c.Query("session_id")
-//	client := types.NewWsClient(ws)
-//	// 删除旧的连接
-//	h.clients.Delete(sessionId)
-//	h.clients.Put(sessionId, client)
-//	logger.Infof("New websocket connected, IP: %s", c.ClientIP())
-//}
-//
-//type sdNotifyData struct {
-//	TaskId    string
-//	ImageName string
-//	ImageData string
-//	Progress  int
-//	Seed      string
-//	Success   bool
-//	Message   string
-//}
-//
-//func (h *SdJobHandler) Notify(c *gin.Context) {
-//	token := c.GetHeader("Authorization")
-//	if token != h.App.Config.ExtConfig.Token {
-//		resp.NotAuth(c)
-//		return
-//	}
-//	var data sdNotifyData
-//	if err := c.ShouldBindJSON(&data); err != nil || data.TaskId == "" {
-//		resp.ERROR(c, types.InvalidArgs)
-//		return
-//	}
-//	logger.Debugf("收到 MidJourney 回调请求：%+v", data)
-//
-//	h.lock.Lock()
-//	defer h.lock.Unlock()
-//
-//	err, finished := h.notifyHandler(c, data)
-//	if err != nil {
-//		resp.ERROR(c, err.Error())
-//		return
-//	}
-//
-//	// 解除任务锁定
-//	if finished && (data.Progress == 100) {
-//		h.redis.Del(c, service.MjRunningJobKey)
-//	}
-//	resp.SUCCESS(c)
-//
-//}
-//
-//func (h *SdJobHandler) notifyHandler(c *gin.Context, data sdNotifyData) (error, bool) {
-//	taskString, err := h.redis.Get(c, service.MjRunningJobKey).Result()
-//	if err != nil { // 过期任务，丢弃
-//		logger.Warn("任务已过期：", err)
-//		return nil, true
-//	}
-//
-//	var task types.SdTask
-//	err = utils.JsonDecode(taskString, &task)
-//	if err != nil { // 非标准任务，丢弃
-//		logger.Warn("任务解析失败：", err)
-//		return nil, false
-//	}
-//
-//	var job model.SdJob
-//	res := h.db.Where("id = ?", task.Id).First(&job)
-//	if res.Error != nil {
-//		logger.Warn("非法任务：", res.Error)
-//		return nil, false
-//	}
-//	job.Params = utils.JsonEncode(task.Params)
-//	job.ReferenceId = data.ImageData
-//	job.Progress = data.Progress
-//	job.Prompt = data.Prompt
-//	job.Hash = data.Image.Hash
-//
-//	// 任务完成，将最终的图片下载下来
-//	if data.Progress == 100 {
-//		imgURL, err := h.uploaderManager.GetUploadHandler().PutImg(data.Image.URL)
-//		if err != nil {
-//			logger.Error("error with download img: ", err.Error())
-//			return err, false
-//		}
-//		job.ImgURL = imgURL
-//	} else {
-//		// 临时图片直接保存，访问的时候使用代理进行转发
-//		job.ImgURL = data.Image.URL
-//	}
-//	res = h.db.Updates(&job)
-//	if res.Error != nil {
-//		logger.Error("error with update job: ", res.Error)
-//		return res.Error, false
-//	}
-//
-//	var jobVo vo.MidJourneyJob
-//	err := utils.CopyObject(job, &jobVo)
-//	if err == nil {
-//		if data.Progress < 100 {
-//			image, err := utils.DownloadImage(jobVo.ImgURL, h.App.Config.ProxyURL)
-//			if err == nil {
-//				jobVo.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
-//			}
-//		}
-//
-//		// 推送任务到前端
-//		client := h.clients.Get(task.SessionId)
-//		if client != nil {
-//			utils.ReplyChunkMessage(client, jobVo)
-//		}
-//	}
-//
-//	// 更新用户剩余绘图次数
-//	if data.Progress == 100 {
-//		h.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
-//	}
-//
-//	return nil, true
-//}
-//
-//func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
-//	user, err := utils.GetLoginUser(c, h.db)
-//	if err != nil {
-//		resp.NotAuth(c)
-//		return false
-//	}
-//
-//	if user.ImgCalls <= 0 {
-//		resp.ERROR(c, "您的绘图次数不足，请联系管理员充值！")
-//		return false
-//	}
-//
-//	return true
-//
-//}
-//
-//// Image 创建一个绘画任务
-//func (h *SdJobHandler) Image(c *gin.Context) {
-//	var data struct {
-//		SessionId string  `json:"session_id"`
-//		Prompt    string  `json:"prompt"`
-//		Rate      string  `json:"rate"`
-//		Model     string  `json:"model"`
-//		Chaos     int     `json:"chaos"`
-//		Raw       bool    `json:"raw"`
-//		Seed      int64   `json:"seed"`
-//		Stylize   int     `json:"stylize"`
-//		Img       string  `json:"img"`
-//		Weight    float32 `json:"weight"`
-//	}
-//	if err := c.ShouldBindJSON(&data); err != nil {
-//		resp.ERROR(c, types.InvalidArgs)
-//		return
-//	}
-//	if !h.checkLimits(c) {
-//		return
-//	}
-//
-//	var prompt = data.Prompt
-//	if data.Rate != "" && !strings.Contains(prompt, "--ar") {
-//		prompt += " --ar " + data.Rate
-//	}
-//	if data.Seed > 0 && !strings.Contains(prompt, "--seed") {
-//		prompt += fmt.Sprintf(" --seed %d", data.Seed)
-//	}
-//	if data.Stylize > 0 && !strings.Contains(prompt, "--s") && !strings.Contains(prompt, "--stylize") {
-//		prompt += fmt.Sprintf(" --s %d", data.Stylize)
-//	}
-//	if data.Chaos > 0 && !strings.Contains(prompt, "--c") && !strings.Contains(prompt, "--chaos") {
-//		prompt += fmt.Sprintf(" --c %d", data.Chaos)
-//	}
-//	if data.Img != "" {
-//		prompt = fmt.Sprintf("%s %s", data.Img, prompt)
-//		if data.Weight > 0 {
-//			prompt += fmt.Sprintf(" --iw %f", data.Weight)
-//		}
-//	}
-//	if data.Raw {
-//		prompt += " --style raw"
-//	}
-//	if data.Model != "" && !strings.Contains(prompt, "--v") && !strings.Contains(prompt, "--niji") {
-//		prompt += data.Model
-//	}
-//
-//	idValue, _ := c.Get(types.LoginUserID)
-//	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
-//	job := model.MidJourneyJob{
-//		Type:      service.Image.String(),
-//		UserId:    userId,
-//		Progress:  0,
-//		Prompt:    prompt,
-//		CreatedAt: time.Now(),
-//	}
-//	if res := h.db.Create(&job); res.Error != nil {
-//		resp.ERROR(c, "添加任务失败："+res.Error.Error())
-//		return
-//	}
-//
-//	h.mjService.PushTask(service.MjTask{
-//		Id:        int(job.Id),
-//		SessionId: data.SessionId,
-//		Src:       service.TaskSrcImg,
-//		Type:      service.Image,
-//		Prompt:    prompt,
-//		UserId:    userId,
-//	})
-//
-//	var jobVo vo.MidJourneyJob
-//	err := utils.CopyObject(job, &jobVo)
-//	if err == nil {
-//		// 推送任务到前端
-//		client := h.clients.Get(data.SessionId)
-//		if client != nil {
-//			utils.ReplyChunkMessage(client, jobVo)
-//		}
-//	}
-//	resp.SUCCESS(c)
-//}
-//
-//// JobList 获取 MJ 任务列表
-//func (h *SdJobHandler) JobList(c *gin.Context) {
-//	status := h.GetInt(c, "status", 0)
-//	var items []model.MidJourneyJob
-//	var res *gorm.DB
-//	userId, _ := c.Get(types.LoginUserID)
-//	if status == 1 {
-//		res = h.db.Where("user_id = ? AND progress = 100", userId).Order("id DESC").Find(&items)
-//	} else {
-//		res = h.db.Where("user_id = ? AND progress < 100", userId).Order("id ASC").Find(&items)
-//	}
-//	if res.Error != nil {
-//		resp.ERROR(c, types.NoData)
-//		return
-//	}
-//
-//	var jobs = make([]vo.MidJourneyJob, 0)
-//	for _, item := range items {
-//		var job vo.MidJourneyJob
-//		err := utils.CopyObject(item, &job)
-//		if err != nil {
-//			continue
-//		}
-//		if item.Progress < 100 {
-//			// 30 分钟还没完成的任务直接删除
-//			if time.Now().Sub(item.CreatedAt) > time.Minute*30 {
-//				h.db.Delete(&item)
-//				continue
-//			}
-//			if item.ImgURL != "" { // 正在运行中任务使用代理访问图片
-//				image, err := utils.DownloadImage(item.ImgURL, h.App.Config.ProxyURL)
-//				if err == nil {
-//					job.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
-//				}
-//			}
-//		}
-//		jobs = append(jobs, job)
-//	}
-//	resp.SUCCESS(c, jobs)
-//}
+import (
+	"chatplus/core"
+	"chatplus/core/types"
+	"chatplus/service/sd"
+	"chatplus/store/model"
+	"chatplus/store/vo"
+	"chatplus/utils"
+	"chatplus/utils/resp"
+	"encoding/base64"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/go-redis/redis/v8"
+	"github.com/gorilla/websocket"
+	"gorm.io/gorm"
+	"net/http"
+	"time"
+)
+
+type SdJobHandler struct {
+	BaseHandler
+	redis   *redis.Client
+	db      *gorm.DB
+	service *sd.Service
+}
+
+func NewSdJobHandler(app *core.AppServer, redisCli *redis.Client, db *gorm.DB, service *sd.Service) *SdJobHandler {
+	h := SdJobHandler{
+		redis:   redisCli,
+		db:      db,
+		service: service,
+	}
+	h.App = app
+	return &h
+}
+
+// Client WebSocket 客户端，用于通知任务状态变更
+func (h *SdJobHandler) Client(c *gin.Context) {
+	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error(err)
+		return
+	}
+
+	sessionId := c.Query("session_id")
+	client := types.NewWsClient(ws)
+	// 删除旧的连接
+	h.service.Clients.Put(sessionId, client)
+	logger.Infof("New websocket connected, IP: %s", c.ClientIP())
+}
+
+func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
+	user, err := utils.GetLoginUser(c, h.db)
+	if err != nil {
+		resp.NotAuth(c)
+		return false
+	}
+
+	if user.ImgCalls <= 0 {
+		resp.ERROR(c, "您的绘图次数不足，请联系管理员充值！")
+		return false
+	}
+
+	return true
+
+}
+
+// Image 创建一个绘画任务
+func (h *SdJobHandler) Image(c *gin.Context) {
+	if !h.App.Config.SdConfig.Enabled {
+		resp.ERROR(c, "Stable Diffusion service is disabled")
+		return
+	}
+
+	if !h.checkLimits(c) {
+		return
+	}
+
+	var data struct {
+		SessionId string `json:"session_id"`
+		types.SdTaskParams
+	}
+	if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
+		resp.ERROR(c, types.InvalidArgs)
+		return
+	}
+
+	if data.Width <= 0 {
+		data.Width = 512
+	}
+	if data.Height <= 0 {
+		data.Height = 512
+	}
+	if data.CfgScale <= 0 {
+		data.CfgScale = 7
+	}
+	if data.Seed == 0 {
+		data.Seed = -1
+	}
+	if data.Steps <= 0 {
+		data.Steps = 20
+	}
+	if data.Sampler == "" {
+		data.Sampler = "Euler a"
+	}
+	idValue, _ := c.Get(types.LoginUserID)
+	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
+	params := types.SdTaskParams{
+		TaskId:         fmt.Sprintf("task(%s)", utils.RandString(15)),
+		Prompt:         data.Prompt,
+		NegativePrompt: data.NegativePrompt,
+		Steps:          data.Steps,
+		Sampler:        data.Sampler,
+		FaceFix:        data.FaceFix,
+		CfgScale:       data.CfgScale,
+		Seed:           data.Seed,
+		Height:         data.Height,
+		Width:          data.Width,
+		HdFix:          data.HdFix,
+		HdRedrawRate:   data.HdRedrawRate,
+		HdScale:        data.HdScale,
+		HdScaleAlg:     data.HdScaleAlg,
+		HdSteps:        data.HdSteps,
+	}
+	job := model.SdJob{
+		UserId:    userId,
+		Type:      types.TaskImage.String(),
+		TaskId:    params.TaskId,
+		Params:    utils.JsonEncode(params),
+		Prompt:    data.Prompt,
+		Progress:  0,
+		Started:   false,
+		CreatedAt: time.Now(),
+	}
+	res := h.db.Create(&job)
+	if res.Error != nil {
+		resp.ERROR(c, "error with save job: "+res.Error.Error())
+		return
+	}
+
+	h.service.PushTask(types.SdTask{
+		Id:        int(job.Id),
+		SessionId: data.SessionId,
+		Src:       types.TaskSrcImg,
+		Type:      types.TaskImage,
+		Prompt:    data.Prompt,
+		Params:    params,
+		UserId:    userId,
+	})
+	var jobVo vo.SdJob
+	err := utils.CopyObject(job, &jobVo)
+	if err == nil {
+		// 推送任务到前端
+		client := h.service.Clients.Get(data.SessionId)
+		if client != nil {
+			utils.ReplyChunkMessage(client, jobVo)
+		}
+	}
+	resp.SUCCESS(c)
+}
+
+// JobList 获取 MJ 任务列表
+func (h *SdJobHandler) JobList(c *gin.Context) {
+	status := h.GetInt(c, "status", 0)
+	var items []model.SdJob
+	var res *gorm.DB
+	userId, _ := c.Get(types.LoginUserID)
+	if status == 1 {
+		res = h.db.Where("user_id = ? AND progress = 100", userId).Order("id DESC").Find(&items)
+	} else {
+		res = h.db.Where("user_id = ? AND progress < 100", userId).Order("id ASC").Find(&items)
+	}
+	if res.Error != nil {
+		resp.ERROR(c, types.NoData)
+		return
+	}
+
+	var jobs = make([]vo.SdJob, 0)
+	for _, item := range items {
+		var job vo.SdJob
+		err := utils.CopyObject(item, &job)
+		if err != nil {
+			continue
+		}
+		if item.Progress < 100 {
+			// 30 分钟还没完成的任务直接删除
+			if time.Now().Sub(item.CreatedAt) > time.Minute*30 {
+				h.db.Delete(&item)
+				continue
+			}
+			if item.ImgURL != "" { // 正在运行中任务使用代理访问图片
+				image, err := utils.DownloadImage(item.ImgURL, h.App.Config.ProxyURL)
+				if err == nil {
+					job.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
+				}
+			}
+		}
+		jobs = append(jobs, job)
+	}
+	resp.SUCCESS(c, jobs)
+}